{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BI4R8LJu0NTW"
      },
      "outputs": [],
      "source": [
        "#@title License\n",
        "license = \"\"\"\n",
        "Copyright 2022 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not\n",
        "use this file except in compliance with the License. You may obtain a copy of\n",
        "the License at https://www.apache.org/licenses/LICENSE-2.0. Unless required by\n",
        "applicable law or agreed to in writing, software distributed under the License\n",
        "is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n",
        "KIND, either express or implied. See the License for the specific language\n",
        "governing permissions and limitations under the License.\n",
        "\"\"\"\n",
        "print(license)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "smVKjDIBg5l-"
      },
      "source": [
        "# Dramatron\n",
        "\n",
        "**Piotr Mirowski, Kory Mathewson, Juliette Love, Jaylen Pittman**, based on an original prototype by **Richard Evans**.\n",
        "\n",
        "DeepMind, 2022\n",
        "\n",
        "Dramatron is a script writing tool that leverages large language models. It uses hierarchical generation to ensure long range consistency across the entire script. The work was published as a pre-print in Mirowski, Mathewson et al. (2022) \"[Co-Writing Screenplays and Theatre Scripts with Language Models: An Evaluation by Industry Professionals](https://arxiv.org/abs/2209.14958)\".\n",
        "\n",
        "![picture](https://github.com/deepmind/dramatron/raw/main/dramatron_logo.png)\n",
        "\n",
        "## Citing this work (paper and code):\n",
        "```bibtex\n",
        "@article{mirowski2022cowriting,\n",
        "  title={Co-Writing Screenplays and Theatre Scripts with Language Models: An Evaluation by Industry Professionals},\n",
        "  author={Mirowski, Piotr and Mathewson, Kory W and Pittman, Jaylen and Evans, Richard},\n",
        "  journal={arXiv preprint arXiv:2209.14958},\n",
        "  year={2022}\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SiVzyJEZHyMs"
      },
      "source": [
        "# Setup\n",
        "\n",
        "Run this set of cells to install and import libraries, define the code, set the parameters of Dramatron and Perspective API (moderation, optional) and the prompt sets."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Google PaLM 2 only: Install dependencies (Vertex AI) and restart the colab runtime\n",
        "\n",
        "#@markdown The Python SDK for the Palm 2 API and the Gemini API, is contained in the [`google-generativeai`](https://pypi.org/project/google-generativeai/) package. Install the dependency using pip:\n",
        "!pip install -q -U google-generativeai\n",
        "# !pip install google-cloud-aiplatform --upgrade --user"
      ],
      "metadata": {
        "id": "Qf5r_vaxIr_M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7WiSfJBhfnSX"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "#@markdown Run this cell to import the required Python libraries.\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "import datetime\n",
        "import difflib\n",
        "import json\n",
        "import re\n",
        "import requests\n",
        "import sys\n",
        "import time\n",
        "from typing import Dict, List, NamedTuple, Optional, Union\n",
        "\n",
        "import collections\n",
        "from google.colab import files\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)\n",
        "\n",
        "print('Imports imported successfully!')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-RsP1ZNd3yH6"
      },
      "outputs": [],
      "source": [
        "#@title Dramatron hyperparameters { run: \"auto\" }\n",
        "\n",
        "#@markdown Run this cell to set the model hyperparameters. The cell will auto-run if you change them. We tested the model with these default values.\n",
        "\n",
        "#@markdown Default seed for generation (default: 1)\n",
        "DEFAULT_SEED =  1  #@param {type:\"integer\"}\n",
        "#@markdown Sampling top-p probability (default: 0.9) and temperature (default 1.0)\n",
        "SAMPLING_PROB =  0.9  #@param {type:\"slider\", min:0.8, max:1.0, step:0.01}\n",
        "SAMPLING_TEMP =  1.  #@param {type:\"slider\", min:0.8, max:1.0, step:0.01}\n",
        "#@markdown Max length for the generated title, place description and others, in tokens (defaults: 64, 128 and 511 respectively)\n",
        "SAMPLE_LENGTH_TITLE = 64 #@param [64, 128, 256, 511]\n",
        "SAMPLE_LENGTH_PLACE = 128 #@param [128, 256, 511]\n",
        "SAMPLE_LENGTH = 511 #@param [128, 256, 511]\n",
        "#@markdown Max lengths during repeated sampling, in case `<end>` is not found (default: 2048)\n",
        "MAX_PARAGRAPH_LENGTH_CHARACTERS = 1024 #@param [511, 1024, 2048, 4096]\n",
        "MAX_PARAGRAPH_LENGTH_SCENES = 1024 #@param [511, 1024, 2048, 4096]\n",
        "MAX_PARAGRAPH_LENGTH = 1024 #@param [511, 1024, 2048, 4096]\n",
        "#@markdown Unavailable API: max number of retries before giving up (default: 10)\n",
        "MAX_RETRIES = 10 #@param {type:\"slider\", min:1, max:20, step:1}\n",
        "#@markdown Loop detection: max number of repetitions before resampling, and number of attempts to get out of the loop (default: 3)\n",
        "MAX_NUM_REPETITIONS = 3 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "MAX_NUM_ATTEMPTS_GET_OUT_OF_LOOP = 3 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "\n",
        "print('Dramatron hyperparameters set.')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vWwxotb_A1w_"
      },
      "outputs": [],
      "source": [
        "#@title Dramatron code\n",
        "\n",
        "#@markdown This cell contains Dramatron's code. You need to run it only once.\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# Script markers\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "# END_MARKER = '<end>'\n",
        "# STOP_MARKER = '<stop>'\n",
        "# CHARACTER_MARKER = '<character>'\n",
        "# DESCRIPTION_MARKER = '<description>'\n",
        "# SCENES_MARKER = '<scenes>'\n",
        "# DIALOG_MARKER = '<dialog>'\n",
        "# TITLE_ELEMENT = 'Title: '\n",
        "# CHARACTERS_ELEMENT = 'Characters: '\n",
        "# DESCRIPTION_ELEMENT = 'Description: '\n",
        "# PLACE_ELEMENT = 'Place: '\n",
        "# PLOT_ELEMENT = 'Plot element: '\n",
        "# PREVIOUS_ELEMENT = 'Previous beat: '\n",
        "# SUMMARY_ELEMENT = 'Summary: '\n",
        "# BEAT_ELEMENT = 'Beat: '\n",
        "END_MARKER = '**END**'\n",
        "STOP_MARKER = '\\n'\n",
        "CHARACTER_MARKER = '**Character:** '\n",
        "DESCRIPTION_MARKER = '**Description:** '\n",
        "SCENES_MARKER = '**Scenes:**'\n",
        "DIALOG_MARKER = '**Dialog:**'\n",
        "LOGLINE_MARKER = \"**Logline:** \"\n",
        "TITLE_ELEMENT = 'Title: '\n",
        "CHARACTERS_ELEMENT = 'Characters: '\n",
        "DESCRIPTION_ELEMENT = 'Description: '\n",
        "PLACE_ELEMENT = 'Place: '\n",
        "PLOT_ELEMENT = 'Plot element: '\n",
        "PREVIOUS_ELEMENT = 'Previous beat: '\n",
        "SUMMARY_ELEMENT = 'Summary: '\n",
        "BEAT_ELEMENT = 'Beat: '\n",
        "LOGLINE_ELEMENT = \"Logline: \"\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# Dramatron script entities\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "\n",
        "class Title(NamedTuple):\n",
        "  \"\"\"Title class.\"\"\"\n",
        "\n",
        "  title: str\n",
        "\n",
        "  @classmethod\n",
        "  def from_string(cls, text: str):\n",
        "    title = extract_elements(text, TITLE_ELEMENT, END_MARKER)[0]\n",
        "    return cls(title)\n",
        "\n",
        "  def to_string(self):\n",
        "    s = ''\n",
        "    s += TITLE_ELEMENT + self.title\n",
        "    s += END_MARKER\n",
        "    return s\n",
        "\n",
        "\n",
        "def get_title(title: Title) -> str:\n",
        "  return title.title\n",
        "\n",
        "\n",
        "class Character(NamedTuple):\n",
        "  \"\"\"Character class.\"\"\"\n",
        "\n",
        "  # Name of the character.\n",
        "  name: str\n",
        "\n",
        "  # A single sentence describing the character.\n",
        "  description: str\n",
        "\n",
        "  @classmethod\n",
        "  def from_string(cls, text: str):\n",
        "    elements = text.split(DESCRIPTION_MARKER)\n",
        "    if len(elements) == 2:\n",
        "      name = elements[0].strip()\n",
        "      description = elements[1].strip()\n",
        "      return cls(name, description)\n",
        "    else:\n",
        "      return None\n",
        "\n",
        "\n",
        "def get_character_description(character: Character) -> str:\n",
        "  return character.description\n",
        "\n",
        "\n",
        "class Characters(NamedTuple):\n",
        "  \"\"\"Characters class, containing main characters and their descriptions.\"\"\"\n",
        "\n",
        "  # A dictionary of character descriptions.\n",
        "  character_descriptions: Dict[str, str]\n",
        "\n",
        "  @classmethod\n",
        "  def from_string(cls, text: str):\n",
        "    \"\"\"Parses the characters from the generated text.\"\"\"\n",
        "    text = text.strip()\n",
        "\n",
        "    # Extracts the character descriptions.\n",
        "    character_descriptions = {}\n",
        "    elements = extract_elements(text, CHARACTER_MARKER, STOP_MARKER)\n",
        "    for text_character in elements:\n",
        "      character = Character.from_string(text_character)\n",
        "      if character is not None:\n",
        "        character_descriptions[character.name] = character.description\n",
        "    return cls(character_descriptions)\n",
        "\n",
        "  def to_string(self):\n",
        "    s = '\\n'\n",
        "    for name, description in self.character_descriptions.items():\n",
        "      s += '\\n' + CHARACTER_MARKER + ' ' + name + ' ' + DESCRIPTION_MARKER + ' '\n",
        "      s += description + ' ' + STOP_MARKER + '\\n'\n",
        "    s += END_MARKER\n",
        "    return s\n",
        "\n",
        "\n",
        "def get_character_descriptions(characters: Characters) -> Dict[str, str]:\n",
        "  return characters.character_descriptions\n",
        "\n",
        "\n",
        "class Scene(NamedTuple):\n",
        "  \"\"\"Scene class.\"\"\"\n",
        "\n",
        "  # The name of the place where the scene unfolds.\n",
        "  place: str\n",
        "\n",
        "  # Name of the plot element (e.g., Beginning, Middle, Conclusion).\n",
        "  plot_element: str\n",
        "\n",
        "  # A short description of action/story/dramatic event occuring in the scene.\n",
        "  beat: str\n",
        "\n",
        "  def to_string(self):\n",
        "    s = PLACE_ELEMENT + ' ' + self.place + '\\n'\n",
        "    s += PLOT_ELEMENT + ' ' + self.plot_element + '\\n'\n",
        "    s += BEAT_ELEMENT + ' ' + self.beat + '\\n'\n",
        "    return s\n",
        "\n",
        "\n",
        "class Scenes(NamedTuple):\n",
        "  \"\"\"Scenes class.\"\"\"\n",
        "\n",
        "  # A list of scenes, with place, characters, plot element and beat.\n",
        "  scenes: List[Scene]\n",
        "\n",
        "  @classmethod\n",
        "  def from_string(cls, text: str):\n",
        "    \"\"\"Parse scenes from generated scenes_text.\"\"\"\n",
        "\n",
        "    places = extract_elements(text, PLACE_ELEMENT, PLOT_ELEMENT)\n",
        "    plot_elements = extract_elements(text, PLOT_ELEMENT, BEAT_ELEMENT)\n",
        "    beats = extract_elements(text, BEAT_ELEMENT, '\\n')\n",
        "\n",
        "    # Get the number of complete scenes.\n",
        "    num_complete_scenes = min([len(places), len(plot_elements), len(beats)])\n",
        "    scenes = []\n",
        "    for i in range(num_complete_scenes):\n",
        "      scenes.append(\n",
        "          Scene(Place.format_name(places[i]), plot_elements[i], beats[i]))\n",
        "    scenes = cls(scenes)\n",
        "    return scenes\n",
        "\n",
        "  def to_string(self):\n",
        "    s = ''\n",
        "    for scene in self.scenes:\n",
        "      s += '\\n' + scene.to_string()\n",
        "    s += END_MARKER\n",
        "    return s\n",
        "\n",
        "  def num_places(self):\n",
        "    return len(set([scene.place for scene in self.scenes]))\n",
        "\n",
        "  def num_scenes(self) -> int:\n",
        "    return len(self.scenes)\n",
        "\n",
        "\n",
        "class Place(NamedTuple):\n",
        "  \"\"\"Place class.\"\"\"\n",
        "\n",
        "  # Place name.\n",
        "  name: str\n",
        "\n",
        "  # Place description.\n",
        "  description: str\n",
        "\n",
        "  @classmethod\n",
        "  def format_name(cls, name: str):\n",
        "    if name.find('.') == -1:\n",
        "      name = name + '.'\n",
        "    return name\n",
        "\n",
        "  @classmethod\n",
        "  def from_string(cls, place_name: str, place_text: str):\n",
        "    place_text += END_MARKER\n",
        "    description = extract_elements(place_text, DESCRIPTION_ELEMENT, END_MARKER)\n",
        "    return cls(place_name, description[0])\n",
        "\n",
        "  @classmethod\n",
        "  def format_prefix(cls, name):\n",
        "    s = PLACE_ELEMENT + name + '\\n' + DESCRIPTION_ELEMENT\n",
        "    return s\n",
        "\n",
        "  def to_string(self):\n",
        "    s = self.format_prefix(self.name) + self.description + '\\n\\n'\n",
        "    return s\n",
        "\n",
        "\n",
        "def get_place_description(place: Place):\n",
        "  return place.description\n",
        "\n",
        "\n",
        "class Story(NamedTuple):\n",
        "  \"\"\"Story class.\"\"\"\n",
        "\n",
        "  # A storyline is a single sentence summary of the whole plot.\n",
        "  storyline: str\n",
        "\n",
        "  # A title for the story.\n",
        "  title: str\n",
        "\n",
        "  # Map from character names to full descriptions.\n",
        "  character_descriptions: Dict[str, str]\n",
        "\n",
        "  # Map from place names to full descriptions.\n",
        "  place_descriptions: Dict[str, Place]\n",
        "\n",
        "  # List of scenes.\n",
        "  scenes: Scenes\n",
        "\n",
        "  # List of dialogs, one for each scene.\n",
        "  dialogs: List[str]\n",
        "\n",
        "\n",
        "def extract_elements(text: str, begin: str, end: str) -> List[str]:\n",
        "  \"\"\"Extracts elements from a text string given string and ending markers.\"\"\"\n",
        "\n",
        "  results = []\n",
        "  start = 0\n",
        "  while True:\n",
        "    start = text.find(begin, start)\n",
        "    if start == -1:\n",
        "      return results\n",
        "    finish = text.find(end, start)\n",
        "    if finish == -1:\n",
        "      return results\n",
        "    results.append(text[start + len(begin):finish].strip())\n",
        "    start = finish + len(end)\n",
        "\n",
        "\n",
        "def strip_remove_end(text: str) -> str:\n",
        "  text = text.strip()\n",
        "  end_marker_stripped = END_MARKER.strip()\n",
        "  if text.endswith(end_marker_stripped):\n",
        "    text = text[:-len(end_marker_stripped)]\n",
        "  return text\n",
        "\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# Rendering of generated stories\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "\n",
        "def render_story(story: Story) -> str:\n",
        "  \"\"\"Render the story in fountain format.\"\"\"\n",
        "\n",
        "  lines = []\n",
        "  lines.append(f'Title: {story.title}')\n",
        "  lines.append('Author: Co-written by ________ and Dramatron')\n",
        "  lines.append(\n",
        "      'Dramatron was developed by Piotr Mirowski and Kory W. Mathewson, '\n",
        "      'with additional contributions by Juliette Love and Jaylen Pittman, '\n",
        "      'and is based on a prototype by Richard Evans.')\n",
        "  lines.append('Dramatron relies on user-provided language models.')\n",
        "  lines.append('')\n",
        "  lines.append('====')\n",
        "  lines.append('')\n",
        "\n",
        "  lines.append(f'The script is based on the storyline:\\n{story.storyline}')\n",
        "  lines.append('')\n",
        "  if story.character_descriptions is not None:\n",
        "    for name, description in story.character_descriptions.items():\n",
        "      lines.append(f'{name}: {description}')\n",
        "      lines.append('')\n",
        "\n",
        "  # For each scene, render scene information.\n",
        "  if story.scenes is not None:\n",
        "    scenes = story.scenes.scenes\n",
        "    for i, scene in enumerate(scenes):\n",
        "      lines.append(f'Scene {i+1}')\n",
        "      lines.append(f'{PLACE_ELEMENT}{scene.place}')\n",
        "      lines.append(f'{PLOT_ELEMENT}{scene.plot_element}')\n",
        "      lines.append(f'{BEAT_ELEMENT}{scene.beat}')\n",
        "      lines.append('')\n",
        "  else:\n",
        "    scenes = []\n",
        "\n",
        "  lines.append('====')\n",
        "  lines.append('')\n",
        "\n",
        "  # For each scene, render the scene's place description, characters and dialog.\n",
        "  for i, scene in enumerate(scenes):\n",
        "\n",
        "    # Output the places and place descriptions.\n",
        "    lines.append(f'INT/EXT. {scene.place} - Scene {i+1}')\n",
        "    place_descriptions = story.place_descriptions\n",
        "    if (not place_appears_earlier(scene.place, story, i) and\n",
        "        place_descriptions is not None and scene.place in place_descriptions):\n",
        "      lines.append('')\n",
        "      lines.append(get_place_description(place_descriptions[scene.place]))\n",
        "\n",
        "    # Output the characters and descriptions.\n",
        "    lines.append('')\n",
        "    for c in story.character_descriptions.keys():\n",
        "      if c in scene.beat and not character_appears_earlier(c, story, i):\n",
        "        lines.append(story.character_descriptions[c])\n",
        "\n",
        "    # Output the dialog.\n",
        "    if story.dialogs is not None and len(story.dialogs) > i:\n",
        "      lines.append('')\n",
        "      lines_dialog = strip_remove_end(str(story.dialogs[i]))\n",
        "      lines.append(lines_dialog)\n",
        "      lines.append('')\n",
        "      lines.append('')\n",
        "\n",
        "  return '\\n'.join(lines)\n",
        "\n",
        "\n",
        "def place_appears_earlier(place: str, story: Story, index: int) -> bool:\n",
        "  \"\"\"Return True if the place appears earlier in the story.\"\"\"\n",
        "\n",
        "  for i in range(index):\n",
        "    scene = story.scenes.scenes[i]\n",
        "    if scene.place == place:\n",
        "      return True\n",
        "  return False\n",
        "\n",
        "\n",
        "def character_appears_earlier(character: str, story: Story, index: int) -> bool:\n",
        "  \"\"\"Return True if the character appears earlier in the story.\"\"\"\n",
        "\n",
        "  for i in range(index):\n",
        "    scene = story.scenes.scenes[i]\n",
        "    if character in scene.beat:\n",
        "      return True\n",
        "  return False\n",
        "\n",
        "\n",
        "def render_prompts(prompts):\n",
        "  \"\"\"Render the prompts.\"\"\"\n",
        "\n",
        "  def _format_prompt(prompt, name):\n",
        "    prompt_str = '=' * 80 + '\\n'\n",
        "    prompt_str += 'PROMPT (' + name + ')\\n'\n",
        "    prompt_str += '=' * 80 + '\\n\\n'\n",
        "    prompt_str += str(prompt) + '\\n\\n'\n",
        "    return prompt_str\n",
        "\n",
        "  prompts_str = _format_prompt(prompts['title'], 'title')\n",
        "  prompts_str += _format_prompt(prompts['characters'], 'characters')\n",
        "  prompts_str += _format_prompt(prompts['scenes'], 'scenes')\n",
        "  places = prompts['places']\n",
        "  if places is not None:\n",
        "    for k, prompt in enumerate(places):\n",
        "      prompts_str += _format_prompt(prompt, 'place ' + str(k + 1))\n",
        "  dialogs = prompts['dialogs']\n",
        "  if dialogs is not None:\n",
        "    for k, prompt in enumerate(dialogs):\n",
        "      prompts_str += _format_prompt(prompt, 'dialog ' + str(k + 1))\n",
        "  return prompts_str\n",
        "\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# Language API definition\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "_MAX_RETRIES = 10\n",
        "_TIMEOUT = 120.0\n",
        "\n",
        "\n",
        "class LanguageResponse(NamedTuple):\n",
        "  prompt: str\n",
        "  prompt_length: int\n",
        "  text: str\n",
        "  text_length: int\n",
        "\n",
        "\n",
        "class LanguageAPI:\n",
        "  \"\"\"Language model wrapper.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               sample_length: int,\n",
        "               model: Optional[str] = None,\n",
        "               model_param: Optional[str] = None,\n",
        "               config_sampling: Optional[dict] = None,\n",
        "               seed: Optional[int] = None,\n",
        "               max_retries: int = _MAX_RETRIES,\n",
        "               timeout: float = _TIMEOUT):\n",
        "    \"\"\"Initializer.\n",
        "\n",
        "    Args:\n",
        "      sample_length: Length of text to sample from model.\n",
        "      model: The model name to correct to. An error will be raised if it does\n",
        "        not exist.\n",
        "      model_param: Model parameter.\n",
        "      config_sampling: Sampleing parameters.\n",
        "      seed: Random seed for sampling.\n",
        "      max_retries: Maximum number of retries for the remote API.\n",
        "      timeout: Maximum waiting timeout\n",
        "    \"\"\"\n",
        "    self._sample_length = sample_length\n",
        "    self._model = model\n",
        "    self._model_param = model_param\n",
        "    self._config_sampling = config_sampling\n",
        "    self._seed = seed\n",
        "    self._max_retries = max_retries\n",
        "    self._timeout = timeout\n",
        "\n",
        "  @property\n",
        "  def default_sample_length(self):\n",
        "    return self._sample_length\n",
        "\n",
        "  @property\n",
        "  def model(self):\n",
        "    return self._model\n",
        "\n",
        "  @property\n",
        "  def model_param(self):\n",
        "    return self._model_param\n",
        "\n",
        "  @property\n",
        "  def model_metadata(self):\n",
        "    return None\n",
        "\n",
        "  @property\n",
        "  def seed(self):\n",
        "    return self._seed\n",
        "\n",
        "  @property\n",
        "  def config_sampling(self):\n",
        "    return self._config_sampling\n",
        "\n",
        "  def sample(self,\n",
        "             prompt: str,\n",
        "             sample_length: Optional[int] = None,\n",
        "             seed: Optional[int] = None,\n",
        "             num_samples: int = 1):\n",
        "    \"\"\"Sample model with provided prompt, optional sample_length and seed.\"\"\"\n",
        "    raise NotImplementedError('sample method not implemented in generic class')\n",
        "\n",
        "\n",
        "class FilterAPI:\n",
        "  \"\"\"Filter model wrapper.\"\"\"\n",
        "\n",
        "  def validate(self, text: str):\n",
        "    raise NotImplementedError('validate not implemented in generic class')\n",
        "\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# Dramatron Generator\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "\n",
        "def generate_text(generation_prompt: str,\n",
        "                  client: LanguageAPI,\n",
        "                  filter: Optional[FilterAPI] = None,\n",
        "                  sample_length: Optional[int] = None,\n",
        "                  max_paragraph_length: int = MAX_PARAGRAPH_LENGTH,\n",
        "                  seed: Optional[int] = None,\n",
        "                  num_samples: int = 1,\n",
        "                  max_num_repetitions: Optional[int] = None) -> str:\n",
        "  \"\"\"Generate text using the generation prompt.\"\"\"\n",
        "\n",
        "  # To prevent lengthy generation loops, we cap the number of calls to the API.\n",
        "  if sample_length is None:\n",
        "    sample_length = client.default_sample_length\n",
        "  max_num_calls = int(max_paragraph_length / sample_length) + 1\n",
        "  num_calls = 0\n",
        "\n",
        "  result = ''\n",
        "  while True:\n",
        "    prompt = generation_prompt + result\n",
        "    success, current_seed = False, seed\n",
        "    while success is False:\n",
        "      t0 = time.time()\n",
        "      responses = client.sample(\n",
        "          prompt=prompt,\n",
        "          sample_length=sample_length,\n",
        "          seed=current_seed,\n",
        "          num_samples=num_samples)\n",
        "      t1 = time.time()\n",
        "      # Get the first result from the list of responses\n",
        "      response = responses[0]\n",
        "      if filter is not None and not filter.validate(response.text):\n",
        "        return 'Content was filtered out.' + END_MARKER\n",
        "      if max_num_repetitions:\n",
        "        success = not detect_loop(\n",
        "            response.text, max_num_repetitions=max_num_repetitions)\n",
        "        if not success:\n",
        "          current_seed += 1\n",
        "          if current_seed > (seed + MAX_NUM_ATTEMPTS_GET_OUT_OF_LOOP):\n",
        "            success = True\n",
        "          else:\n",
        "            continue\n",
        "      else:\n",
        "        success = True\n",
        "\n",
        "    result = result + response.text\n",
        "    num_calls += 1\n",
        "\n",
        "    # Attempt to find the END_MARKER\n",
        "    index = result.find(END_MARKER)\n",
        "    if index != -1:\n",
        "      return result[:index] + END_MARKER\n",
        "\n",
        "    # Attempt to find the start of a new example\n",
        "    index = result.find('Example ')\n",
        "    if index != -1:\n",
        "      return result[:index] + END_MARKER\n",
        "\n",
        "    if max_paragraph_length is not None and len(result) > max_paragraph_length:\n",
        "      return result + END_MARKER\n",
        "    if num_calls >= max_num_calls:\n",
        "      return result + END_MARKER\n",
        "\n",
        "  return result\n",
        "\n",
        "\n",
        "def generate_text_no_loop(generation_prompt: str,\n",
        "                          client: LanguageAPI,\n",
        "                          filter: Optional[FilterAPI] = None,\n",
        "                          sample_length: Optional[int] = None,\n",
        "                          max_paragraph_length: int = MAX_PARAGRAPH_LENGTH,\n",
        "                          seed: Optional[int] = None,\n",
        "                          num_samples: int = 1) -> str:\n",
        "  \"\"\"Generate text using the generation prompt, without any loop.\"\"\"\n",
        "  return generate_text(\n",
        "      generation_prompt=generation_prompt,\n",
        "      client=client,\n",
        "      filter=filter,\n",
        "      sample_length=sample_length,\n",
        "      max_paragraph_length=sample_length,\n",
        "      seed=seed,\n",
        "      max_num_repetitions=None,\n",
        "      num_samples=num_samples)\n",
        "\n",
        "\n",
        "def generate_title(storyline: str,\n",
        "                   prefixes: Dict[str, str],\n",
        "                   client: LanguageAPI,\n",
        "                   filter: Optional[FilterAPI] = None,\n",
        "                   seed: Optional[int] = None,\n",
        "                   num_samples: int = 1):\n",
        "  \"\"\"Generate a title given a storyline, and client.\"\"\"\n",
        "\n",
        "  # Combine the prompt and storyline as a helpful generation prefix\n",
        "  titles_prefix = prefixes['TITLES_PROMPT'] + storyline + ' ' + TITLE_ELEMENT\n",
        "  title_text = generate_text_no_loop(\n",
        "      generation_prompt=titles_prefix,\n",
        "      client=client,\n",
        "      filter=filter,\n",
        "      sample_length=SAMPLE_LENGTH_TITLE,\n",
        "      seed=seed,\n",
        "      num_samples=num_samples)\n",
        "  title = Title.from_string(TITLE_ELEMENT + title_text)\n",
        "  return (title, titles_prefix)\n",
        "\n",
        "\n",
        "def generate_characters(\n",
        "    storyline: str,\n",
        "    prefixes: Dict[str, str],\n",
        "    client: LanguageAPI,\n",
        "    filter: Optional[FilterAPI] = None,\n",
        "    seed: Optional[int] = None,\n",
        "    max_paragraph_length: int = (MAX_PARAGRAPH_LENGTH_CHARACTERS),\n",
        "    num_samples: int = 1):\n",
        "  \"\"\"Generate characters given a storyline, prompt, and client.\"\"\"\n",
        "\n",
        "  # Combine the prompt and storyline as a helpful generation prefix\n",
        "  characters_prefix = prefixes['CHARACTERS_PROMPT'] + storyline\n",
        "  characters_text = generate_text(\n",
        "      generation_prompt=characters_prefix,\n",
        "      client=client,\n",
        "      filter=filter,\n",
        "      seed=seed,\n",
        "      max_paragraph_length=max_paragraph_length,\n",
        "      num_samples=num_samples)\n",
        "  characters = Characters.from_string(characters_text)\n",
        "\n",
        "  return (characters, characters_prefix)\n",
        "\n",
        "\n",
        "def generate_scenes(storyline: str,\n",
        "                    character_descriptions: Dict[str, str],\n",
        "                    prefixes: Dict[str, str],\n",
        "                    client: LanguageAPI,\n",
        "                    filter: Optional[FilterAPI] = None,\n",
        "                    seed: Optional[int] = None,\n",
        "                    max_paragraph_length: int = (MAX_PARAGRAPH_LENGTH_SCENES),\n",
        "                    num_samples: int = 1):\n",
        "  \"\"\"Generate scenes given storyline, prompt, main characters, and client.\"\"\"\n",
        "\n",
        "  scenes_prefix = prefixes['SCENE_PROMPT'] + storyline + '\\n'\n",
        "  for name in character_descriptions:\n",
        "    scenes_prefix += character_descriptions[name] + '\\n'\n",
        "  scenes_prefix += '\\n' + SCENES_MARKER\n",
        "  scenes_text = generate_text(\n",
        "      generation_prompt=scenes_prefix,\n",
        "      client=client,\n",
        "      filter=filter,\n",
        "      seed=seed,\n",
        "      max_paragraph_length=max_paragraph_length,\n",
        "      num_samples=num_samples)\n",
        "  scenes = Scenes.from_string(scenes_text)\n",
        "\n",
        "  return (scenes, scenes_prefix)\n",
        "\n",
        "\n",
        "def generate_place_descriptions(storyline: str,\n",
        "                                scenes: Scenes,\n",
        "                                prefixes: Dict[str, str],\n",
        "                                client: LanguageAPI,\n",
        "                                filter: Optional[FilterAPI] = None,\n",
        "                                seed: Optional[int] = None,\n",
        "                                num_samples: int = 1):\n",
        "  \"\"\"Generate a place description given a scene object and a client.\"\"\"\n",
        "\n",
        "  place_descriptions = {}\n",
        "\n",
        "  # Get unique place names from the scenes.\n",
        "  unique_place_names = set([scene.place for scene in scenes.scenes])\n",
        "\n",
        "  # Build a unique place prefix prompt.\n",
        "  place_prefix = prefixes['SETTING_PROMPT'] + storyline + '\\n'\n",
        "\n",
        "  # Build a list of place descriptions for each place\n",
        "  place_prefixes = []\n",
        "  for place_name in unique_place_names:\n",
        "    place_suffix = Place.format_prefix(place_name)\n",
        "    place_text = generate_text(\n",
        "        generation_prompt=place_prefix + place_suffix,\n",
        "        client=client,\n",
        "        filter=filter,\n",
        "        sample_length=SAMPLE_LENGTH_PLACE,\n",
        "        seed=seed,\n",
        "        num_samples=num_samples)\n",
        "    place_text = place_suffix + place_text\n",
        "    place_descriptions[place_name] = Place.from_string(place_name, place_text)\n",
        "    place_prefixes.append(place_prefix + place_suffix)\n",
        "\n",
        "  return (place_descriptions, place_prefixes)\n",
        "\n",
        "\n",
        "def prefix_summary(storyline: str,\n",
        "                   scenes: List[Scene],\n",
        "                   concatenate_scenes_in_summary: bool = False) -> str:\n",
        "  \"\"\"Assemble the summary part of the dialog prefix.\"\"\"\n",
        "\n",
        "  summary = SUMMARY_ELEMENT + storyline + '\\n'\n",
        "  if len(scenes) > 1:\n",
        "    summary += PREVIOUS_ELEMENT + scenes[len(scenes) - 2].beat + '\\n'\n",
        "  return summary\n",
        "\n",
        "\n",
        "def detect_loop(text: str, max_num_repetitions: int = MAX_NUM_REPETITIONS):\n",
        "  \"\"\"Detect loops in generated text.\"\"\"\n",
        "\n",
        "  blocks = text.split('\\n\\n')\n",
        "  num_unique_blocks = collections.Counter(blocks)\n",
        "  for block in blocks:\n",
        "    num_repetitions = num_unique_blocks[block]\n",
        "    if num_repetitions > max_num_repetitions:\n",
        "      print(f'Detected {num_repetitions} repetitions of block:\\n{block}')\n",
        "      return True\n",
        "  return False\n",
        "\n",
        "\n",
        "def generate_dialog(storyline: str,\n",
        "                    scenes: List[Scene],\n",
        "                    character_descriptions: Dict[str, str],\n",
        "                    place_descriptions: Dict[str, Place],\n",
        "                    prefixes: Dict[str, str],\n",
        "                    max_paragraph_length: int,\n",
        "                    client: LanguageAPI,\n",
        "                    filter: Optional[FilterAPI] = None,\n",
        "                    max_num_repetitions: Optional[int] = None,\n",
        "                    seed: Optional[int] = None,\n",
        "                    num_samples: int = 1):\n",
        "  \"\"\"Generate dialog given a scene object and a client.\"\"\"\n",
        "\n",
        "  scene = scenes[-1]\n",
        "\n",
        "  place_t = PLACE_ELEMENT + scene.place + '\\n'\n",
        "  if scene.place in place_descriptions:\n",
        "    place_description = place_descriptions[scene.place]\n",
        "    if place_description:\n",
        "      place_t += DESCRIPTION_ELEMENT + place_description.description\n",
        "      place_t += '\\n'\n",
        "\n",
        "  # Build the characters information for the scene\n",
        "  characters_t = ''\n",
        "  if character_descriptions:\n",
        "    characters_t += CHARACTERS_ELEMENT\n",
        "    for name in character_descriptions:\n",
        "      if name in scene.beat:\n",
        "        characters_t += character_descriptions[name] + '\\n'\n",
        "\n",
        "  plot_element_t = PLOT_ELEMENT + scene.plot_element + '\\n'\n",
        "\n",
        "  summary_t = prefix_summary(\n",
        "      storyline, scenes, concatenate_scenes_in_summary=False)\n",
        "\n",
        "  beat_t = BEAT_ELEMENT + scene.beat + '\\n'\n",
        "\n",
        "  dialog_prefix = (\n",
        "      prefixes['DIALOG_PROMPT'] + place_t + characters_t + plot_element_t +\n",
        "      summary_t + beat_t)\n",
        "  dialog_prefix += '\\n' + DIALOG_MARKER + '\\n'\n",
        "\n",
        "  dialog = generate_text(\n",
        "      generation_prompt=dialog_prefix,\n",
        "      client=client,\n",
        "      filter=filter,\n",
        "      seed=seed,\n",
        "      max_paragraph_length=max_paragraph_length,\n",
        "      max_num_repetitions=max_num_repetitions,\n",
        "      num_samples=num_samples)\n",
        "\n",
        "  return (dialog, dialog_prefix)\n",
        "\n",
        "\n",
        "def diff_prompt_change_str(prompt_before: str, prompt_after: str) -> str:\n",
        "  \"\"\"Return a text diff on prompt sets `prompt_before` and `prompt_after`.\"\"\"\n",
        "\n",
        "  # For the current element, compare prompts line by line.\n",
        "  res = difflib.unified_diff(\n",
        "      prompt_before.split('\\n'), prompt_after.split('\\n'))\n",
        "  diff = ''\n",
        "  for line in res:\n",
        "    line = line.strip()\n",
        "    if line != '---' and line != '+++' and not line.startswith('@@'):\n",
        "      if len(line) > 1 and (line.startswith('+') or line.startswith('-')):\n",
        "        diff += line + '\\n'\n",
        "  if diff.endswith('\\n'):\n",
        "    diff = diff[:-1]\n",
        "  return diff\n",
        "\n",
        "\n",
        "def diff_prompt_change_list(prompt_before: List[str],\n",
        "                            prompt_after: List[str]) -> str:\n",
        "  \"\"\"Return a text diff on prompt sets `prompt_before` and `prompt_after`.\"\"\"\n",
        "\n",
        "  # Handle deletions and insertions.\n",
        "  len_before = len(prompt_before)\n",
        "  len_after = len(prompt_after)\n",
        "  if len_before > len_after:\n",
        "    return 'Deleted element'\n",
        "  if len_before < len_after:\n",
        "    return 'Added new element'\n",
        "\n",
        "  diffs = [\n",
        "      diff_prompt_change_str(a, b)\n",
        "      for (a, b) in zip(prompt_before, prompt_after)\n",
        "  ]\n",
        "  return '\\n'.join([diff for diff in diffs if len(diff) > 0])\n",
        "\n",
        "\n",
        "def diff_prompt_change_scenes(prompt_before: List[Scene],\n",
        "                              prompt_after: List[Scene]) -> str:\n",
        "  \"\"\"Return a text diff on prompt sets `prompt_before` and `prompt_after`.\"\"\"\n",
        "\n",
        "  # Handle deletions and insertions.\n",
        "  len_before = len(prompt_before)\n",
        "  len_after = len(prompt_after)\n",
        "  if len_before > len_after:\n",
        "    return 'Deleted element'\n",
        "  if len_before < len_after:\n",
        "    return 'Added new element'\n",
        "\n",
        "  diffs = [\n",
        "      diff_prompt_change_list([a.place, a.plot_element, a.beat],\n",
        "                              [b.place, b.plot_element, b.beat])\n",
        "      for (a, b) in zip(prompt_before, prompt_after)\n",
        "  ]\n",
        "  return '\\n'.join([diff for diff in diffs if len(diff) > 0])\n",
        "\n",
        "\n",
        "def diff_prompt_change_dict(prompt_before: Dict[str, str],\n",
        "                            prompt_after: Dict[str, str]) -> str:\n",
        "  \"\"\"Return a text diff on prompt sets `prompt_before` and `prompt_after`.\"\"\"\n",
        "\n",
        "  # Loop over the keys in the prompts to compare them one by one.\n",
        "  keys_before = sorted(prompt_before.keys())\n",
        "  keys_after = sorted(prompt_after.keys())\n",
        "  diffs = [\n",
        "      diff_prompt_change_str(a, b) for (a, b) in zip(keys_before, keys_after)\n",
        "  ]\n",
        "  diff_keys = '\\n'.join([diff for diff in diffs if len(diff) > 0])\n",
        "  # Loop over the values in the prompts to compare them one by one.\n",
        "  values_before = sorted(prompt_before.values())\n",
        "  values_after = sorted(prompt_after.values())\n",
        "  diffs = [\n",
        "      diff_prompt_change_str(a, b)\n",
        "      for (a, b) in zip(values_before, values_after)\n",
        "  ]\n",
        "  diff_values = '\\n'.join([diff for diff in diffs if len(diff) > 0])\n",
        "  return diff_keys + diff_values\n",
        "\n",
        "\n",
        "class StoryGenerator:\n",
        "  \"\"\"Generate a story from the provided storyline, using the client provided.\"\"\"\n",
        "\n",
        "  level_names = ('storyline', 'title', 'characters', 'scenes', 'places',\n",
        "                 'dialogs')\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      storyline: str,\n",
        "      prefixes: Dict[str, str],\n",
        "      max_paragraph_length: int = 1024,\n",
        "      max_paragraph_length_characters: int = (MAX_PARAGRAPH_LENGTH_CHARACTERS),\n",
        "      max_paragraph_length_scenes: int = (MAX_PARAGRAPH_LENGTH_SCENES),\n",
        "      num_samples: int = 1,\n",
        "      client: Optional[LanguageAPI] = None,\n",
        "      filter: Optional[FilterAPI] = None):\n",
        "    self._prefixes = prefixes\n",
        "    self._max_paragraph_length = max_paragraph_length\n",
        "    self._max_paragraph_length_characters = max_paragraph_length_characters\n",
        "    self._max_paragraph_length_scenes = max_paragraph_length_scenes\n",
        "    self._num_samples = num_samples\n",
        "    self._client = client\n",
        "    self._filter = filter\n",
        "\n",
        "    # Prompts and outputs of the hierarchical generator are organised in levels.\n",
        "    self.prompts = {\n",
        "        'title': '',\n",
        "        'characters': '',\n",
        "        'scenes': '',\n",
        "        'places': {\n",
        "            '': ''\n",
        "        },\n",
        "        'dialogs': ['']\n",
        "    }\n",
        "    self._title = Title('')\n",
        "    self._characters = Characters({'': ''})\n",
        "    self._scenes = Scenes([Scene('', '', '')])\n",
        "    self._places = {'': Place('', '')}\n",
        "    self._dialogs = ['']\n",
        "\n",
        "    # History of interventions.\n",
        "    self.interventions = {}\n",
        "    self._set_storyline(storyline)\n",
        "\n",
        "  def _set_storyline(self, storyline: str):\n",
        "    \"\"\"Set storyline and initialise the outputs of the generator.\"\"\"\n",
        "    self._level = 0\n",
        "\n",
        "    # Add period to the end of the storyline, unless there is already one there.\n",
        "    if storyline.find('.') == -1:\n",
        "      storyline = storyline + '.'\n",
        "    self._storyline = storyline\n",
        "\n",
        "    # Keep track of each storyline intervention.\n",
        "    timestamp = time.time()\n",
        "    self.interventions[timestamp] = 'STORYLINE\\n' + storyline\n",
        "\n",
        "  @property\n",
        "  def seed(self):\n",
        "    return self._client.seed\n",
        "\n",
        "  @property\n",
        "  def title(self) -> Title:\n",
        "    \"\"\"Return the title.\"\"\"\n",
        "    return self._title\n",
        "\n",
        "  @property\n",
        "  def characters(self) -> Characters:\n",
        "    \"\"\"Return the characters.\"\"\"\n",
        "    return self._characters\n",
        "\n",
        "  @property\n",
        "  def scenes(self) -> Scenes:\n",
        "    \"\"\"Return the title.\"\"\"\n",
        "    return self._scenes\n",
        "\n",
        "  @property\n",
        "  def places(self) -> Dict[str, Place]:\n",
        "    \"\"\"Return the places.\"\"\"\n",
        "    return self._places\n",
        "\n",
        "  @property\n",
        "  def dialogs(self) -> List[str]:\n",
        "    \"\"\"Return the dialogs.\"\"\"\n",
        "    return self._dialogs\n",
        "\n",
        "  def title_str(self) -> str:\n",
        "    \"\"\"Return the title as a string.\"\"\"\n",
        "    return self._title.title\n",
        "\n",
        "  def num_scenes(self) -> int:\n",
        "    \"\"\"Return the number of scenes.\"\"\"\n",
        "    return self._scenes.num_scenes()\n",
        "\n",
        "  def step(self,\n",
        "           level: Optional[int] = None,\n",
        "           seed: Optional[int] = None,\n",
        "           idx: Optional[int] = None) -> bool:\n",
        "    \"\"\"Step down a level in the hierarchical generation of a story.\"\"\"\n",
        "\n",
        "    # Move to the next level of hierarchical generation.\n",
        "    if level is None:\n",
        "      level = self._level\n",
        "    if level < 0 or level >= len(self.level_names):\n",
        "      raise ValueError('Invalid level encountered on step.')\n",
        "    level += 1\n",
        "    self._level = level\n",
        "\n",
        "    # Keep track of each step intervention.\n",
        "    timestamp = time.time()\n",
        "    self.interventions[timestamp] = 'STEP ' + str(level) + '\\n'\n",
        "\n",
        "    if level == 1:\n",
        "      # Step 1: Generate title given a storyline.\n",
        "      (title, titles_prefix) = generate_title(\n",
        "          storyline=self._storyline,\n",
        "          prefixes=self._prefixes,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          num_samples=self._num_samples,\n",
        "          seed=seed)\n",
        "      self._title = title\n",
        "      self.prompts['title'] = titles_prefix\n",
        "      self.interventions[timestamp] += title.to_string()\n",
        "      success = len(title.title) > 0\n",
        "      return success\n",
        "\n",
        "    if level == 2:\n",
        "      # Step 2: Generate characters given a storyline.\n",
        "      (characters, character_prompts) = generate_characters(\n",
        "          storyline=self._storyline,\n",
        "          prefixes=self._prefixes,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          num_samples=self._num_samples,\n",
        "          max_paragraph_length=self._max_paragraph_length_characters,\n",
        "          seed=seed)\n",
        "      self._characters = characters\n",
        "      self.prompts['characters'] = character_prompts\n",
        "      self.interventions[timestamp] += characters.to_string()\n",
        "      success = len(characters.character_descriptions) > 0\n",
        "      return success\n",
        "\n",
        "    if level == 3:\n",
        "      # Step 3: Generate sequence of scenes given a storyline and characters.\n",
        "      characters = self._characters\n",
        "      (scenes, scene_prompts) = generate_scenes(\n",
        "          storyline=self._storyline,\n",
        "          character_descriptions=get_character_descriptions(characters),\n",
        "          prefixes=self._prefixes,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          num_samples=self._num_samples,\n",
        "          max_paragraph_length=self._max_paragraph_length_scenes,\n",
        "          seed=seed)\n",
        "      self._scenes = scenes\n",
        "      self.prompts['scenes'] = scene_prompts\n",
        "      self.interventions[timestamp] += scenes.to_string()\n",
        "      success = len(scenes.scenes) > 0\n",
        "      return success\n",
        "\n",
        "    if level == 4:\n",
        "      # Step 4: For each scene, generate place descriptions given place name.\n",
        "      scenes = self._scenes\n",
        "      (place_descriptions, place_prompts) = generate_place_descriptions(\n",
        "          storyline=self._storyline,\n",
        "          scenes=scenes,\n",
        "          prefixes=self._prefixes,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          num_samples=self._num_samples,\n",
        "          seed=seed)\n",
        "      self._places = place_descriptions\n",
        "      self.prompts['places'] = place_prompts\n",
        "      for place_name in place_descriptions:\n",
        "        place = place_descriptions[place_name]\n",
        "        if place:\n",
        "          self.interventions[timestamp] += place.to_string()\n",
        "      num_places = scenes.num_places()\n",
        "      success = (len(place_descriptions) == num_places) and num_places > 0\n",
        "      return success\n",
        "\n",
        "    if level == 5:\n",
        "      # Step 5: For each scene, generate dialog from scene information.\n",
        "      title = self._title\n",
        "      characters = self._characters\n",
        "      scenes = self._scenes\n",
        "      place_descriptions = self._places\n",
        "      if idx is None:\n",
        "        (dialogs, dialog_prompts) = zip(*[\n",
        "            generate_dialog(\n",
        "                storyline=self._storyline,\n",
        "                scenes=scenes.scenes[:(k + 1)],\n",
        "                character_descriptions=(characters.character_descriptions),\n",
        "                place_descriptions=place_descriptions,\n",
        "                prefixes=self._prefixes,\n",
        "                max_paragraph_length=self._max_paragraph_length,\n",
        "                max_num_repetitions=MAX_NUM_REPETITIONS,\n",
        "                client=self._client,\n",
        "                filter=self._filter,\n",
        "                num_samples=self._num_samples,\n",
        "                seed=seed) for k in range(len(scenes.scenes))\n",
        "        ])\n",
        "      else:\n",
        "        num_scenes = self._scenes.num_scenes()\n",
        "        while len(self._dialogs) < num_scenes:\n",
        "          self._dialogs.append('')\n",
        "        while len(self.prompts['dialogs']) < num_scenes:\n",
        "          self.prompts['dialogs'].append('')\n",
        "        if idx >= num_scenes or idx < 0:\n",
        "          raise ValueError('Invalid scene index.')\n",
        "        dialogs = self._dialogs\n",
        "        dialog_prompts = self.prompts['dialogs']\n",
        "        dialogs[idx], dialog_prompts[idx] = generate_dialog(\n",
        "            storyline=self._storyline,\n",
        "            scenes=scenes.scenes[:(idx + 1)],\n",
        "            character_descriptions=(characters.character_descriptions),\n",
        "            place_descriptions=place_descriptions,\n",
        "            prefixes=self._prefixes,\n",
        "            max_paragraph_length=self._max_paragraph_length,\n",
        "            max_num_repetitions=MAX_NUM_REPETITIONS,\n",
        "            client=self._client,\n",
        "            filter=self._filter,\n",
        "            num_samples=self._num_samples,\n",
        "            seed=seed)\n",
        "      self._dialogs = dialogs\n",
        "      self.prompts['dialogs'] = dialog_prompts\n",
        "      for dialog in dialogs:\n",
        "        self.interventions[timestamp] += str(dialog)\n",
        "      return True\n",
        "\n",
        "  def get_story(self):\n",
        "    if self._characters is not None:\n",
        "      character_descriptions = get_character_descriptions(self._characters)\n",
        "    else:\n",
        "      character_descriptions = None\n",
        "    return Story(\n",
        "        storyline=self._storyline,\n",
        "        title=self._title.title,\n",
        "        character_descriptions=character_descriptions,\n",
        "        place_descriptions=self._places,\n",
        "        scenes=self._scenes,\n",
        "        dialogs=self._dialogs)\n",
        "\n",
        "  def rewrite(self, text, level=0, entity=None):\n",
        "    if level < 0 or level >= len(self.level_names):\n",
        "      raise ValueError('Invalid level encountered on step.')\n",
        "    prompt_diff = None\n",
        "\n",
        "    if level == 0:\n",
        "      # Step 0: Rewrite the storyline and begin new story.\n",
        "      prompt_diff = diff_prompt_change_str(self._storyline, text)\n",
        "      self._set_storyline(text)\n",
        "\n",
        "    if level == 1:\n",
        "      # Step 1: Rewrite the title.\n",
        "      title = Title.from_string(text)\n",
        "      prompt_diff = diff_prompt_change_str(self._title.title, title.title)\n",
        "      self._title = title\n",
        "\n",
        "    if level == 2:\n",
        "      # Step 2: Rewrite the characters.\n",
        "      characters = Characters.from_string(text)\n",
        "      prompt_diff = diff_prompt_change_dict(\n",
        "          self._characters.character_descriptions,\n",
        "          characters.character_descriptions)\n",
        "      self._characters = characters\n",
        "\n",
        "    if level == 3:\n",
        "      # Step 3: Rewrite the sequence of scenes.\n",
        "      scenes = Scenes.from_string(text)\n",
        "      prompt_diff = diff_prompt_change_scenes(self._scenes.scenes,\n",
        "                                              scenes.scenes)\n",
        "      self._scenes = scenes\n",
        "\n",
        "    if level == 4:\n",
        "      # Step 4: For a given place, rewrite its place description.\n",
        "      place_descriptions = self._places\n",
        "      if entity in place_descriptions:\n",
        "        place_prefix = Place.format_prefix(entity)\n",
        "        text = place_prefix + text\n",
        "        place = Place.from_string(entity, text)\n",
        "        prompt_diff = diff_prompt_change_str(self._places[entity].name,\n",
        "                                             place.name)\n",
        "        prompt_diff += '\\n' + diff_prompt_change_str(\n",
        "            self._places[entity].description, place.description)\n",
        "\n",
        "        self._places[entity] = place\n",
        "\n",
        "    if level == 5:\n",
        "      # Step 5: Rewrite the dialog of a given scene.\n",
        "      dialogs = self._dialogs\n",
        "      num_scenes = len(self._scenes.scenes)\n",
        "      if entity >= 0 and entity < num_scenes:\n",
        "        prompt_diff = diff_prompt_change_str(self._dialogs[entity], text)\n",
        "        self._dialogs[entity] = text\n",
        "\n",
        "    # Keep track of each rewrite intervention.\n",
        "    if prompt_diff is not None and len(prompt_diff) > 0:\n",
        "      timestamp = time.time()\n",
        "      self.interventions[timestamp] = 'REWRITE ' + self.level_names[level]\n",
        "      if entity:\n",
        "        self.interventions[timestamp] += ' ' + str(entity)\n",
        "      self.interventions[timestamp] += prompt_diff\n",
        "\n",
        "  def complete(self,\n",
        "               level=0,\n",
        "               seed=None,\n",
        "               entity=None,\n",
        "               sample_length=SAMPLE_LENGTH):\n",
        "    if level < 0 or level >= len(self.level_names):\n",
        "      raise ValueError('Invalid level encountered on step.')\n",
        "    prompt_diff = None\n",
        "\n",
        "    if level == 2:\n",
        "      # Step 2: Complete the characters.\n",
        "      text_characters = self._characters.to_string()\n",
        "      text_characters = strip_remove_end(text_characters)\n",
        "      prompt = self.prompts['characters'] + text_characters\n",
        "      text = generate_text(\n",
        "          generation_prompt=prompt,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          sample_length=sample_length,\n",
        "          max_paragraph_length=sample_length,\n",
        "          seed=seed,\n",
        "          num_samples=1)\n",
        "      new_characters = Characters.from_string(text_characters + text)\n",
        "      prompt_diff = diff_prompt_change_dict(\n",
        "          self._characters.character_descriptions,\n",
        "          new_characters.character_descriptions)\n",
        "      self._characters = new_characters\n",
        "\n",
        "    if level == 3:\n",
        "      # Step 3: Complete the sequence of scenes.\n",
        "      text_scenes = self._scenes.to_string()\n",
        "      text_scenes = strip_remove_end(text_scenes)\n",
        "      prompt = self.prompts['scenes'] + text_scenes\n",
        "      text = generate_text(\n",
        "          generation_prompt=prompt,\n",
        "          client=self._client,\n",
        "          filter=self._filter,\n",
        "          sample_length=sample_length,\n",
        "          max_paragraph_length=sample_length,\n",
        "          seed=seed,\n",
        "          num_samples=1)\n",
        "      new_scenes = Scenes.from_string(text_scenes + text)\n",
        "      prompt_diff = diff_prompt_change_scenes(self._scenes.scenes,\n",
        "                                              new_scenes.scenes)\n",
        "      self._scenes = new_scenes\n",
        "\n",
        "    if level == 5:\n",
        "      # Step 5: Complete the dialog of a given scene.\n",
        "      dialogs = self._dialogs\n",
        "      num_scenes = len(self._scenes.scenes)\n",
        "      while len(self._dialogs) < num_scenes:\n",
        "        self._dialogs.append('')\n",
        "      while len(self.prompts['dialogs']) < num_scenes:\n",
        "        self.prompts['dialogs'].append('')\n",
        "      if entity >= 0 and entity < num_scenes:\n",
        "        prompt = (self.prompts['dialogs'][entity] + self._dialogs[entity])\n",
        "        text = generate_text(\n",
        "            generation_prompt=prompt,\n",
        "            client=self._client,\n",
        "            filter=self._filter,\n",
        "            sample_length=sample_length,\n",
        "            max_paragraph_length=sample_length,\n",
        "            seed=seed,\n",
        "            num_samples=1)\n",
        "        new_dialog = self._dialogs[entity] + text\n",
        "        prompt_diff = diff_prompt_change_str(self._dialogs[entity], new_dialog)\n",
        "        self._dialogs[entity] = new_dialog\n",
        "\n",
        "    # Keep track of each rewrite intervention.\n",
        "    if prompt_diff is not None and len(prompt_diff) > 0:\n",
        "      timestamp = time.time()\n",
        "      self.interventions[timestamp] = 'COMPLETE ' + self.level_names[level]\n",
        "      if entity:\n",
        "        self.interventions[timestamp] += ' ' + str(entity)\n",
        "      self.interventions[timestamp] += prompt_diff\n",
        "\n",
        "\n",
        "# ------------------------------------------------------------------------------\n",
        "# UI\n",
        "# ------------------------------------------------------------------------------\n",
        "\n",
        "\n",
        "class GenerationAction:\n",
        "  NEW = 1\n",
        "  CONTINUE = 2\n",
        "  REWRITE = 3\n",
        "\n",
        "\n",
        "class GenerationHistory:\n",
        "  \"\"\"Custom data structure to handle the history of GenerationAction edits:\n",
        "\n",
        "  NEW, CONTINUE or REWRITE. Consecutive REWRITE edits do not add to history.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    self._items = []\n",
        "    self._actions = []\n",
        "    self._idx = -1\n",
        "    self._locked = False\n",
        "\n",
        "  def _plain_add(self, item, action: GenerationAction):\n",
        "    self._items.append(item)\n",
        "    self._actions.append(action)\n",
        "    self._idx = len(self._items) - 1\n",
        "    return self._idx\n",
        "\n",
        "  def add(self, item, action: GenerationAction):\n",
        "    if len(self._items) == 0 or action != GenerationAction.REWRITE:\n",
        "      return self._plain_add(item, action)\n",
        "    last_action = self._actions[-1]\n",
        "    if last_action != GenerationAction.REWRITE:\n",
        "      return self._plain_add(item, action)\n",
        "    self._items[self._idx] = item\n",
        "    return self._idx\n",
        "\n",
        "  def previous(self):\n",
        "    if len(self._items) == 0:\n",
        "      return None\n",
        "    self._idx = max(self._idx - 1, 0)\n",
        "    return self._items[self._idx]\n",
        "\n",
        "  def next(self):\n",
        "    if len(self._items) == 0:\n",
        "      return None\n",
        "    self._idx = min(self._idx + 1, len(self._items) - 1)\n",
        "    return self._items[self._idx]\n",
        "\n",
        "filter = None\n",
        "\n",
        "print('Dramatron set-up complete.')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F3AtyfN0nWSY",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Custom Language API\n",
        "\n",
        "#@markdown This cell should contain code with your language API.\n",
        "\n",
        "class CustomLanguageAPI(LanguageAPI):\n",
        "  \"\"\"A class wrapping the language model API.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               sample_length: int,\n",
        "               model: Optional[str] = None,\n",
        "               model_param: Optional[str] = None,\n",
        "               config_sampling: Optional[dict] = None,\n",
        "               seed: Optional[int] = None,\n",
        "               max_retries: int = _MAX_RETRIES,\n",
        "               timeout: float = _TIMEOUT):\n",
        "    \"\"\"Initializer.\n",
        "\n",
        "    Args:\n",
        "      sample_length: Length of text to sample from model.\n",
        "      model: The model name to correct to. An error will be raised if it does\n",
        "        not exist.\n",
        "      model_param: Custom language model params.\n",
        "      config_sampling: ConfigDict with parameters.\n",
        "      seed: Random seed for sampling.\n",
        "      max_retries: Maximum number of retries for the remote API.\n",
        "      timeout: Maximum waiting timeout\n",
        "    \"\"\"\n",
        "    raise NotImplementedError('init method not implemented')\n",
        "\n",
        "  @property\n",
        "  def client(self):\n",
        "    return self._client\n",
        "\n",
        "  @property\n",
        "  def model_metadata(self):\n",
        "    return {'engine': self._model,\n",
        "            'model_param': self._model_param,\n",
        "            'max_tokens': self._sample_length}\n",
        "\n",
        "  def sample(self,\n",
        "             prompt: str,\n",
        "             sample_length: Optional[int] = None,\n",
        "             seed: Optional[int] = None,\n",
        "             num_samples: int = 1):\n",
        "    raise NotImplementedError('sample method not implemented')\n",
        "\n",
        "\n",
        "# Create the config.\n",
        "config = {}\n",
        "config['language_api_name'] = None\n",
        "config['model_param'] = None\n",
        "config['model_name'] = None\n",
        "config['max_retries'] = MAX_RETRIES\n",
        "config['sample_length'] = SAMPLE_LENGTH\n",
        "config['max_paragraph_length'] = MAX_PARAGRAPH_LENGTH\n",
        "config['max_paragraph_length_characters'] = MAX_PARAGRAPH_LENGTH_CHARACTERS\n",
        "config['max_paragraph_length_scenes'] = MAX_PARAGRAPH_LENGTH_SCENES\n",
        "config['sampling'] = {}\n",
        "config['sampling']['prob'] = SAMPLING_PROB\n",
        "config['sampling']['temp'] = SAMPLING_TEMP\n",
        "config['sampling']['frequency_penalty'] = 0.23\n",
        "config['sampling']['presence_penalty'] = 0.23\n",
        "config['prefixes'] = {}\n",
        "config['file_dir'] = None\n",
        "\n",
        "print('Config:')\n",
        "for key, value in config.items():\n",
        "  if key != 'prefixes':\n",
        "    print(f'{key}: {value}')\n",
        "\n",
        "try:\n",
        "  client = CustomLanguageAPI(\n",
        "      model_param=config['model_param'],\n",
        "      model=config['model_name'],\n",
        "      seed=DEFAULT_SEED,\n",
        "      sample_length=config['sample_length'],\n",
        "      max_retries=config['max_retries'],\n",
        "      config_sampling=config['sampling'])\n",
        "\n",
        "  print(f'Client model metadata: {client.model_metadata}')\n",
        "except:\n",
        "  print('CustomLanguageAPI not implemented yet')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZhfIe4QBSslN"
      },
      "outputs": [],
      "source": [
        "#@title (Optional) Perspective API { run: \"auto\" }\n",
        "\n",
        "PERSPECTIVE_API_URL = 'https://commentanalyzer.googleapis.com/v1alpha1/comments:analyze'\n",
        "\n",
        "#@markdown The language model may generate offensive text. If you choose, you can provide a [Perspective API](https://perspectiveapi.com/) key, which will hide model outputs that exceed a toxicity threshold and prompt you to regenerate the text.\n",
        "PERSPECTIVE_API_KEY = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Perspective API toxicity thresholds across different attributes (default value: 0.8).\n",
        "PERSPECTIVE_API_TOXICITY = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "PERSPECTIVE_API_SEVERE_TOXICITY = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "PERSPECTIVE_API_IDENTITY_ATTACK = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "PERSPECTIVE_API_INSULT = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "PERSPECTIVE_API_SEXUALLY_EXPLICIT = 0.8 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "PERSPECTIVE_API_THRESHOLDS = {\n",
        "    \"TOXICITY\": PERSPECTIVE_API_TOXICITY,\n",
        "    \"SEVERE_TOXICITY\": PERSPECTIVE_API_SEVERE_TOXICITY,\n",
        "    \"IDENTITY_ATTACK\": PERSPECTIVE_API_IDENTITY_ATTACK,\n",
        "    \"INSULT\": PERSPECTIVE_API_INSULT,\n",
        "    \"SEXUALLY_EXPLICIT\": PERSPECTIVE_API_SEXUALLY_EXPLICIT\n",
        "}\n",
        "\n",
        "\n",
        "class PerspectiveAPI(FilterAPI):\n",
        "  \"\"\"Wraps the Perspective API model.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               key: str,\n",
        "               thresholds: dict):\n",
        "    \"\"\"Initializer:\n",
        "\n",
        "    Args:\n",
        "      key: Perspective API key.\n",
        "      thresholds: Thresholds for attributes.\n",
        "    \"\"\"\n",
        "    self._key = key\n",
        "    self._thresholds = thresholds\n",
        "\n",
        "  def get_scores(self, text: str):\n",
        "    \"\"\"Get the scores from the Perspective API comment analyzer for text.\"\"\"\n",
        "\n",
        "    input_data = {\n",
        "        'comment': {'text': text},\n",
        "        'languages': ['en'],\n",
        "        'requestedAttributes': {attribute:{} for attribute in self._thresholds}\n",
        "    }\n",
        "    response = requests.post(\n",
        "        PERSPECTIVE_API_URL,\n",
        "        params={'key': self._key},\n",
        "        headers={'Content-Type': 'application/json'},\n",
        "        data=json.dumps(input_data)\n",
        "    )\n",
        "    output_data = response.json()\n",
        "    scores = {}\n",
        "    for attribute in self._thresholds:\n",
        "      score = output_data['attributeScores'][attribute]['summaryScore']['value']\n",
        "      scores[attribute] = score\n",
        "    return scores\n",
        "\n",
        "  def validate(self, text: str):\n",
        "    \"\"\"Filter text using scores from the Perspective API.\"\"\"\n",
        "\n",
        "    scores = self.get_scores(text)\n",
        "    return all([scores[attribute] <= self._thresholds[attribute]\n",
        "                for attribute in self._thresholds])\n",
        "\n",
        "if PERSPECTIVE_API_KEY:\n",
        "  filter = PerspectiveAPI(\n",
        "      key=PERSPECTIVE_API_KEY,\n",
        "      thresholds=PERSPECTIVE_API_THRESHOLDS\n",
        "  )\n",
        "  print('PerspectiveAPI configured.')\n",
        "else:\n",
        "  filter = None\n",
        "  print('No Perspective API key provided.')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3OAW4PeE1z7N"
      },
      "source": [
        "## Prompt prefix sets\n",
        "\n",
        "Run these cells once to define the available prompt prefixes. These cells contain 3 sets of prompt prefixes: **Medea**, **Sci-Fi** and **Custom**. The first two sets were used in our study, the third one serves as a placeholder that you can customise."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wzf80E5YAi6g"
      },
      "outputs": [],
      "source": [
        "#@title Medea\n",
        "\n",
        "#@markdown Trigger warning: the script contains sensitive topics.\n",
        "\n",
        "#@markdown Log line: `Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess and the wife of Jason, finds her position in the Greek world threatened as Jason leaves Medea for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as Medea's own two sons, after which she escapes to Athens.`\n",
        "\n",
        "#@markdown Based on Ancient Greek tragedy \"Medea\", by Euripides (431 BC). Text of the play taken verbatim from the translation by E. P. Coleridge (1863 -1936). One edit made to replace `CHORUS` by `WOMEN OF CORINTH`.\n",
        "\n",
        "#@markdown Prompts for Medea written from a summary taken from Spark Notes. Prompts for Antigone (Sophocles), The Bacchae (Euripides), The Frogs (Aristophanes) adapted from Wikipedia.\n",
        "\n",
        "#@markdown To encourage the generation of different locations, Aristotle's Unity of Place is not respected, and location `Outside the Royal Palace` is renamed as `Medea's modest home` as well as `On a winged chariot` (even though these are the same locations in the original tragedy).\n",
        "\n",
        "#@markdown References:\n",
        "\n",
        "#@markdown http://classics.mit.edu/Euripides/medea.pl.txt<br> https://en.wikipedia.org/wiki/Medea_(play)<br> https://www.sparknotes.com/lit/medea/<br> https://www.ancient-literature.com/greece_sophocles_antigone.html<br> https://en.wikipedia.org/wiki/The_Bacchae<br> https://www.ancient-literature.com/greece_aristophanes_frogs.html\n",
        "\n",
        "medea_prefixes = {}\n",
        "medea_prefixes['CHARACTERS_PROMPT'] = \"\"\"\n",
        "Here is an example of a logline and a list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess and the wife of Jason, finds her position in the Greek world threatened as Jason leaves Medea for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as Medea's own two sons, after which she escapes to Athens.\n",
        "\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Medea \"\"\" + DESCRIPTION_MARKER + \"\"\" Medea is the protagonist of the play. A sorceress and a princess, she fled her country and family to live with Jason in Corinth, where they established a family of two children and gained a favorable reputation. Jason has divorced Medea and taken up with a new family.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Jason \"\"\" + DESCRIPTION_MARKER + \"\"\" Jason is considered the play's villain, though his evil stems more from weakness than strength. A former adventurer, Jason abandons his wife, Medea, in order to marry the beautiful young daughter of Creon, King of Corinth, and fuels Medea to a revenge.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Women of Corinth \"\"\" + DESCRIPTION_MARKER + \"\"\" The Women of Corinth are a commentator to the action. They fully sympathizes with Medea's plight, excepting her decision to murder her own children.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Creon \"\"\" + DESCRIPTION_MARKER + \"\"\" Creon is the King of Corinth, banishes Medea from the city.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"The Nurse \"\"\" + DESCRIPTION_MARKER + \"\"\" The Nurse is the caretaker of the house and of the children and serves as Medea's confidant.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and the following logline, complete the list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "medea_prefixes['SCENE_PROMPT'] = \"\"\"\n",
        "Here is an example of a logline, a list of characters, and a list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess and the wife of Jason, finds her position in the Greek world threatened as Jason leaves Medea for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as Medea's own two sons, after which she escapes to Athens.\n",
        "Medea is the protagonist of the play. A sorceress and a princess, she fled her country and family to live with Jason in Corinth, where they established a family of two children and gained a favorable reputation. Jason has divorced Medea and taken up with a new family.\n",
        "Jason can be considered the play's villain, though his evil stems more from weakness than strength. A former adventurer, Jason abandons his wife, Medea, in order to marry the beautiful young daughter of Creon, King of Corinth, and fuels Medea to a revenge.\n",
        "The Women of Corinth serve as a commentator to the action. They fully sympathizes with Medea's plight, excepting her decision to murder her own children.\n",
        "The King of Corinth Creon banishes Medea from the city.\n",
        "The Messenger appears only once in the play to bear tragical news.\n",
        "The Nurse is the caretaker of the house and of the children and serves as Medea's confidant.\n",
        "The Tutor of the children is a very minor character and mainly acts as a messenger.\n",
        "\n",
        "\"\"\" + SCENES_MARKER + \"\"\"\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Medea's modest home.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Exposition.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"The Nurse recounts the chain of events that have turned Medea's world to enmity. The Nurse laments how Jason has abandoned Medea and his own children in order to remarry with the daughter of Creon.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Medea's modest home.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Inciting Incident.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"The Nurse confides in the Tutor amd testifies to the emotional shock Jason's betrayal has sparked in Medea. The Tutor shares the Nurse's sympathy for Medea's plight. Medea's first words are cries of helplessness. Medea wishes for her own death.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Medea's modest home.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Conflict.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"The Women of Corinth address Medea and try to reason with Medea and convince her that suicide would be an overreaction. The Nurse recognizes the gravity of Medea's threat.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Rising Action.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Medea pleads to the Nurse that Jason be made to suffer for the suffering he has inflicted upon her. Creon approaches the house and banishes Medea and her children from Corinth. Medea plans on killing her three antagonists, Creon, his daughter and Jason.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Dilemma.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Jason rebuke Medea for publicly expressing her murderous intentions. Jason defends his choice to remarry. Medea refuses Jason's offers and sends him away to his new bride.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Climax.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"When Jason returns, Medea begins to carry out her ruse. Medea fakes regret and break down in false tears of remorse. Determined, Medea sends her children to offer poisoned gifts to Creon's daughter. Medea's children face impending doom.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Falling Action.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"The Messenger frantically runs towards Medea and warns Medea to escape the city as soon as possible. The Messenger reveals that Medea has been identified as the murderer.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Resolution.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Medea and her two dead children are seated in a chariot drawn by dragons. Jason watches in horror and curses himself for having wed Medea and mourns his tragic losses.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On a winged chariot.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Dénouement.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Medea denies Jason the right to a proper burial of his children. She flees to Athens and divines an unheroic death for Jason.\n",
        "\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and the following logline and list of characters, complete the list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "medea_prefixes['SETTING_PROMPT'] = \"\"\"\n",
        "Here are examples of logline, location, and that location's description.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Ella, a waitress, falls in love with her best friend, Allen, a teacher. The two drift apart when Allen makes new friends from a different social class. Ella turns to food to become a famous chef.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The bar.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The bar is dirty, more than a little run down, with most tables empty. The odor of last night's beer and crushed pretzels on the floor permeates the bar.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Grandma Phyllis’ family reunion with her two grandchildren is crashed by two bikers.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The Lawn in Front of Grandma Phyllis's House.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"A big oak tree dominates the yard. There is an old swing set on the lawn, and a bright white fence all around the grass.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess and the wife of Jason, finds her position in the Greek world threatened as Jason leaves Medea for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as Medea's own two sons, after which she escapes to Athens.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"In mythological Ancient Greece, in front of a modest house in Corinth, on the outskirts of a lavish royal palace where wedding preparations are under way.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the examples above and the following logine and location name, complete location description.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "medea_prefixes['TITLES_PROMPT'] = \"\"\"\n",
        "Examples of alternative, original and descriptive titles for known play and film scripts.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess of the kingdom of Colchis, and the wife of Jason, finds her position in the Greek world threatened as Jason leaves her for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as her own two sons, after which she escapes to Athens.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"A Feminist Tale\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Ancient Greek tragedy that deals with Antigone’s burial of her brother Polynices, in defiance of the laws of Creon and the state, and the tragic repercussions of her act of civil disobedience.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"In My Brother's Name\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\" Greek comedy that tells the story of the god Dionysus (also known to the Greeks as Bacchus) who, despairing of the current state of Athens’ tragedians, travels to Hades with his slave Xanthias to bring Euripides back from the dead.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"Dionysus in Hades\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 4.\n",
        "\"\"\" + LOGLINE_ELEMENT\n",
        "\n",
        "\n",
        "medea_prefixes['DIALOG_PROMPT'] = \"\"\"\n",
        "Here is an example of description and scene dialogue from a classical play.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Outside the Royal Palace.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"Before Medea's house in Corinth, near the royal palace of Creon.\n",
        "\"\"\" + CHARACTERS_ELEMENT + \"\"\"Medea is the protagonist of the play. A sorceress and a princess, she fled her country and family to live with Jason in Corinth, where they established a family of two children and gained a favorable reputation. Jason has divorced Medea and taken up with a new family. Jason can be considered the play's villain, though his evil stems more from weakness than strength. A former adventurer, Jason abandons his wife, Medea, in order to marry the beautiful young daughter of Creon, King of Corinth, and fuels Medea to a revenge. The Messenger appears only once in the play to bear tragical news.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Resolution.\n",
        "\"\"\" + SUMMARY_ELEMENT + \"\"\"Ancient Greek tragedy based upon the myth of Jason and Medea. Medea, a former princess and the wife of Jason, finds her position in the Greek world threatened as Jason leaves Medea for a Greek princess of Corinth. Medea takes vengeance on Jason by murdering his new wife as well as Medea's own two sons, after which she escapes to Athens.\n",
        "\"\"\" + PREVIOUS_ELEMENT + \"\"\"The Messenger frantically warns Medea to escape the city as soon as possible. The Messenger reveals that Medea has been identified as the murderer.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"The palace opens its doors, revealing Medea and the two dead children seated in a chariot drawn by dragons. Jason curses himself for having wed Medea and mourns his tragic losses. Medea denies Jason the right to a proper burial of his children. Medea flees to Athens and divines an unheroic death for Jason.\n",
        "\n",
        "\"\"\" + DIALOG_MARKER + \"\"\"\n",
        "\n",
        "WOMEN OF CORINTH\n",
        "Throw wide the doors and see thy children's murdered corpses.\n",
        "\n",
        "JASON\n",
        "Haste, ye slaves, loose the bolts, undo the fastenings, that\n",
        "I may see the sight of twofold woe, my murdered sons and her, whose\n",
        "blood in vengeance I will shed.  (MEDEA appears above the house, on\n",
        "a chariot drawn by dragons; the children's corpses are beside her.)\n",
        "\n",
        "MEDEA\n",
        "Why shake those doors and attempt to loose their bolts, in\n",
        "quest of the dead and me their murderess? From such toil desist. If\n",
        "thou wouldst aught with me, say on, if so thou wilt; but never shalt\n",
        "thou lay hand on me, so swift the steeds the sun, my father's sire,\n",
        "to me doth give to save me from the hand of my foes.\n",
        "\n",
        "JASON\n",
        "Accursed woman! by gods, by me and all mankind abhorred as\n",
        "never woman was, who hadst the heart to stab thy babes, thou their\n",
        "mother, leaving me undone and childless; this hast thou done and still\n",
        "dost gaze upon the sun and earth after this deed most impious. Curses\n",
        "on thee! now perceive what then I missed in the day I brought thee,\n",
        "fraught with doom, from thy home in a barbarian land to dwell in Hellas,\n",
        "traitress to thy sire and to the land that nurtured thee.\n",
        "Perish, vile sorceress, murderess of\n",
        "thy babes! Whilst I must mourn my luckless fate, for I shall ne'er\n",
        "enjoy my new-found bride, nor shall I have the children, whom I bred\n",
        "and reared, alive to say the last farewell to me; nay, I have lost\n",
        "them.\n",
        "\n",
        "MEDEA\n",
        "To this thy speech I could have made a long reply, but Father\n",
        "Zeus knows well all I have done for thee, and the treatment thou hast\n",
        "given me. Yet thou wert not ordained to scorn my love and lead a life\n",
        "of joy in mockery of me, nor was thy royal bride nor Creon, who gave\n",
        "thee a second wife, to thrust me from this land and rue it not. Wherefore,\n",
        "if thou wilt, call me e'en a lioness, and Scylla, whose home is in\n",
        "the Tyrrhene land; for I in turn have wrung thy heart, as well I might.\n",
        "\n",
        "JASON\n",
        "Thou, too, art grieved thyself, and sharest in my sorrow.\n",
        "\n",
        "MEDEA\n",
        "Be well assured I am; but it relieves my pain to know thou\n",
        "canst not mock at me.\n",
        "\n",
        "JASON\n",
        "O my children, how vile a mother ye have found!\n",
        "\n",
        "MEDEA\n",
        "My sons, your father's feeble lust has been your ruin!\n",
        "\n",
        "JASON\n",
        "'Twas not my hand, at any rate, that slew them.\n",
        "\n",
        "MEDEA\n",
        "No, but thy foul treatment of me, and thy new marriage.\n",
        "\n",
        "JASON\n",
        "Didst think that marriage cause enough to murder them?\n",
        "\n",
        "MEDEA\n",
        "Dost think a woman counts this a trifling injury?\n",
        "\n",
        "JASON\n",
        "So she be self-restrained; but in thy eyes all is evil.\n",
        "\n",
        "MEDEA\n",
        "Thy sons are dead and gone. That will stab thy heart.\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and following description, write the dialogue of the scene.\n",
        "\n",
        "\"\"\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gHhprsguAa9h"
      },
      "outputs": [],
      "source": [
        "#@title Sci-Fi\n",
        "\n",
        "#@markdown Log line for Star Wars: Episode IV: `A science-fiction fantasy about a naive but ambitious farm boy from a backwater desert who discovers powers he never knew he had when he teams up with a feisty princess, a mercenary space pilot and an old wizard warrior to lead a ragtag rebellion against the sinister forces of the evil Galactic Empire.`\n",
        "\n",
        "#@markdown Log line taken from chapter \"Creating the killer log line\" by Bill Lundy, in: Ellis, Sherry, and Laurie Lamson. \"Now Write! Mysteries: Suspense, Crime, Thriller, and Other Mystery Fiction Exercises from Today's Best Writers and Teachers.\" Penguin, 2011.\n",
        "\n",
        "#@markdown Characters are adapted from Star Wars: Episode IV - A New Hope (1977) written and directed by George Lucas, produced by Lucasfilm and distributed by 20th Century Fox.\n",
        "\n",
        "#@markdown Breakdown of Star Wars into a Hero Journey taken from: https://thescriptlab.com/features/screenwriting-101/12309-the-heros-journey-breakdown-star-wars/\n",
        "\n",
        "#@markdown Log line for Plan 9 from Outer Space: `Residents of San Fernando Valley are under attack by flying saucers from outer space. The aliens are extraterrestrials who seek to stop humanity from creating a doomsday weapon that could destroy the universe and unleash the living dead to stalk humans who wander into the cemetery looking for evidence of the UFOs. The hero Jeff, an airline pilot, will face the aliens.`\n",
        "\n",
        "#@markdown The script, plot and logline of \"Plan 9 from Outer Space\" is in public domain, available at:<br> http://www.horrorlair.com/scripts/criswell.txt<br> https://en.wikipedia.org/wiki/Plan_9_from_Outer_Space<br> https://www.rottentomatoes.com/m/plan-9-from-outer-space<br>\n",
        "\n",
        "scifi_prefixes = {}\n",
        "scifi_prefixes['CHARACTERS_PROMPT'] = \"\"\"\n",
        "Here is an example of a logline and a list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"A science-fiction fantasy about a naive but ambitious farm boy from a backwater desert who discovers powers he never knew he had when he teams up with a feisty princess, a mercenary space pilot and an old wizard warrior to lead a ragtag rebellion against the sinister forces of the evil Galactic Empire.\n",
        "\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Luke Skywalker \"\"\" + DESCRIPTION_MARKER + \"\"\"Luke Skywalker is the hero. A naive farm boy, he will discover special powers under the guidance of mentor Ben Kenobi.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Ben Kenobi \"\"\" + DESCRIPTION_MARKER + \"\"\"Ben Kenobi is the mentor figure. A recluse Jedi warrior, he will take Luke Skywalker as apprentice.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Darth Vader \"\"\" + DESCRIPTION_MARKER + \"\"\"Darth Vader is the antagonist. As a commander of the evil Galactic Empire, he controls space station The Death Star.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Princess Leia \"\"\" + DESCRIPTION_MARKER + \"\"\"Princess Leia is a feisty and brave leader of the Rebellion. She holds the plans of the Death Star. She will become Luke's friend.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Han Solo \"\"\" + DESCRIPTION_MARKER + \"\"\"Han Solo is a brash mercenary space pilot of the Millenium Falcon and a friend of Chebacca. He will take Luke on his spaceship.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\"Chewbacca \"\"\" + DESCRIPTION_MARKER + \"\"\"Chewbacca is a furry and trustful monster. He is a friend of Han Solo and a copilot on the Millemium Falcon.\"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and the following logline, complete the list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "scifi_prefixes['SCENE_PROMPT'] = \"\"\"\n",
        "Examples of breakdowns of stories into a Hero's Journey structure.\n",
        "\n",
        "Here is an example of a logline, a list of characters, and a list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"A science-fiction fantasy about a naive but ambitious farm boy from a backwater desert who discovers powers he never knew he had when he teams up with a feisty princess, a mercenary space pilot and an old wizard warrior to lead a ragtag rebellion against the sinister forces of the evil Galactic Empire.\n",
        "Luke Skywalker is the hero. A naive farm boy, he will discover special powers under the guidance of mentor Ben Kenobi.\n",
        "Ben Kenobi is the mentor figure. A recluse Jedi warrior, he will take Luke Skywalker as apprentice.\n",
        "Darth Vader is the antagonist. As a commander of the evil Galactic Empire, he controls space station The Death Star.\n",
        "Princess Leia holds the plans of the Death Star. She is feisty and brave. She will become Luke's friend.\n",
        "Han Solo is a brash mercenary space pilot of the Millenium Falcon and a friend of Chebacca. He will take Luke on his spaceship.\n",
        "Chewbacca is a furry and trustful monster. He is a friend of Han Solo and a copilot on the Millemium Falcon.\n",
        "\n",
        "\"\"\" + SCENES_MARKER + \"\"\"\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"A farm on planet Tatooine.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Ordinary World.\n",
        "Beat: Luke Skywalker is living a normal and humble life as a farm boy on his home planet.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Desert of Tatooine.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Call to Adventure.\n",
        "Beat: Luke is called to his adventure by robot R2-D2 and Ben Kenobi. Luke triggers R2-D2’s message from Princess Leia and is intrigued by her message. When R2-D2 escapes to find Ben Kenobi, Luke follows and is later saved by Kenobi, who goes on to tell Luke about his Jedi heritage. Kenobi suggests that he should come with him.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Ben Kenobi's farm.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Refusal of the Call.\n",
        "Beat: Luke refuses Kenobi, telling him that he can take Kenobi and the droids as far as Mos Eisley Spaceport — but he can’t possibly leave his Aunt and Uncle behind for some space adventure.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"A farm on planet Tatooine.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Crossing the First Threshold.\n",
        "Beat: When Luke discovers that the stormtroopers searching for the droids would track them to his farm, he rushes to warn his Aunt and Uncle, only to discover them dead by the hands of the Empire. When Luke returns to Kenobi, he pledges to go with him to Alderaan and learn the ways of the Force like his father before him.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On spaceship The Millennium Falcon.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Tests, Allies, and Enemies.\n",
        "Beat: After Luke, Kenobi, and the droids hire Han Solo and Chewbacca to transport them onto Alderaan, Kenobi begins Luke’s training in the ways of the Force. Wielding his father’s lightsaber, Kenobi challenges Luke. At first, he can’t do it. But then Kenobi Kenobi Luke him to reach out and trust his feelings. Luke succeeds.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On spaceship The Millennium Falcon.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Approach to the Inmost Cave.\n",
        "Beat: The plan to defeat the Galactic Empire is to bring the Death Star plans to Alderaan so that Princess Leia’s father can take them to the Rebellion. However, when they arrive within the system, the planet is destroyed. They come across the Death Star and are pulled in by a tractor beam, now trapped within the Galactic Empire.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On space station The Death Star.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Ordeal.\n",
        "Beat: As Kenobi goes off to deactivate the tractor beam so they can escape, Luke, Han, and Chewbacca discover that Princess Leia is being held on the Death Star with them. They rescue her and escape to the Millennium Falcon, hoping that Kenobi has successfully deactivated the tractor beam. Kenobi later sacrifices himself as Luke watches Darth Vader strike him down. Luke must now avenge his fallen mentor and carry on his teachings.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On space station The Death Star.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Reward.\n",
        "Beat: Luke has saved the princess and retrieved the Death Star plans. They now have the knowledge to destroy the Galactic Empire’s greatest weapon once and for all.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On spaceship The Millennium Falcon.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Road Back.\n",
        "Beat: Luke, Leia, Han, Chewbacca, and the droids are headed to the hidden Rebellion base with the Death Star plans. They are suddenly pursued by incoming TIE-Fighters, forcing Han and Luke to take action to defend the ship and escape with their lives — and the plans. They race to take the plans to the Rebellion and prepare for battle.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"On fighter ship X-Wing.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Resurrection.\n",
        "Beat: The Rebels — along with Luke as an X-Wing pilot — take on the Death Star. The Rebellion and the Galactic Empire wage war in an epic space battle. Luke is the only X-Wing pilot that was able to get within the trenches of the Death Star. But Darth Vader and his wingmen are in hot pursuit. Just as Darth Vader is about to destroy Luke, Han returns and clears the way for Luke. Luke uses the Force to guide his aiming as he fires upon the sole weak point of the deadly Death Star, destroying it for good.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"At the Rebellion base.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"The Return.\n",
        "Beat: Luke and Han return to the Rebellion base, triumphant, as they receive medals for the heroic journey. There is peace throughout the galaxy — at least for now.\n",
        "\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and the following logline and list of characters, complete the list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "scifi_prefixes['SETTING_PROMPT'] = \"\"\"\n",
        "Here are examples of logline, location, and that location's description.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Morgan adopts a new cat, Misterio, who sets a curse on anyone that pets them.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The Adoption Center.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The Adoption Center is a sad place, especially for an unadopted pet. It is full of walls and walls of cages and cages. Inside of each is an abandoned animal, longing for a home. The lighting is dim, gray, buzzing fluorescent.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"James finds a well in his backyard that is haunted by the ghost of Sam.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The well.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The well is buried under grass and hedges. It is at least twenty feet deep, if not more and it is masoned with stones. It is 150 years old at least. It stinks of stale, standing water, and has vines growing up the sides. It is narrow enough to not be able to fit down if you are a grown adult human.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Mr. Dorbenson finds a book at a garage sale that tells the story of his own life. And it ends in a murder!\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The garage sale.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"It is a garage packed with dusty household goods and antiques. There is a box at the back that says FREE and is full of paper back books.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the examples above and the following logine and location name, complete location description.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "scifi_prefixes['TITLES_PROMPT'] = \"\"\"\n",
        "Examples of alternative, original and descriptive titles for known play and film scripts.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"A science-fiction fantasy about a naive but ambitious farm boy from a backwater desert who discovers powers he never knew he had when he teams up with a feisty princess, a mercenary space pilot and an old wizard warrior to lead a ragtag rebellion against the sinister forces of the evil Galactic Empire.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"The Death Star's Menace\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Residents of San Fernando Valley are under attack by flying saucers from outer space. The aliens are extraterrestrials who seek to stop humanity from creating a doomsday weapon that could destroy the universe and unleash the living dead to stalk humans who wander into the cemetery looking for evidence of the UFOs. The hero Jeff, an airline pilot, will face the aliens.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"The Day The Earth Was Saved By Outer Space.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_ELEMENT\n",
        "\n",
        "\n",
        "scifi_prefixes['DIALOG_PROMPT'] = \"\"\"\n",
        "Here is an example of description and scene dialogue from a modern screenplay.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Cockpit of an airplane.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"Cockpit of a modern passenger airplane, American Flight 812.\n",
        "\"\"\" + CHARACTERS_ELEMENT + \"\"\"Jeff is the hero. A man in his early forties, he tries to stay calm in all circumstance. Jeff is now a airline pilot. Danny, a young airplane pilot in his thirties, is eager to learn but can quickly lose his composture. Danny is enamored of Edith. Edith, an experienced stewardess with a good sense of humour, is trustworthy and dependable. Edith likes to tease Danny.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Crossing the First Threshold.\n",
        "\"\"\" + SUMMARY_ELEMENT + \"\"\"Residents of San Fernando Valley are under attack by flying saucers from outer space. The aliens are extraterrestrials who seek to stop humanity from creating a doomsday weapon that could destroy the universe and unleash the living dead to stalk humans who wander into the cemetery looking for evidence of the UFOs. The hero Jeff, an airline pilot, will face the aliens.\n",
        "\"\"\" + PREVIOUS_ELEMENT + \"\"\"Flight captain Jeff reluctantly leaves his wife Paula to go for a two-day flight.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"At the cockpit, flight captain Jeff is preoccupied by the flying saucer appearances and graveyard incidents in his home town, where he left wis wife Paula. Without success, co-pilot Danny and stewardess Edith try to reassure him.\n",
        "\n",
        "\"\"\" + DIALOG_MARKER + \"\"\"\n",
        "\n",
        "DANNY\n",
        "You're mighty silent this trip, Jeff.\n",
        "\n",
        "JEFF\n",
        "Huh?\n",
        "\n",
        "DANNY\n",
        "You haven't spoken ten words since takeoff.\n",
        "\n",
        "JEFF\n",
        "I guess I'm preoccupied, Danny.\n",
        "\n",
        "DANNY\n",
        "We've got thirty-three passengers back there that have time to be preoccupied.\n",
        "Flying this flybird doesn't give you that opportunity.\n",
        "\n",
        "JEFF\n",
        "I guess you're right, Danny.\n",
        "\n",
        "DANNY\n",
        "Paula?\n",
        "\n",
        "JEFF\n",
        "Yeah.\n",
        "\n",
        "DANNY\n",
        "There's nothing wrong between you two?\n",
        "\n",
        "JEFF\n",
        "Oh no, nothing like that.  Just that I'm worried, she being there alone and\n",
        "those strange things flying over the house and those incidents in the graveyard\n",
        "the past few days. It's just got me worried.\n",
        "\n",
        "DANNY\n",
        "Well, I haven't figured out those crazy skybirds yet but I give you fifty to one\n",
        "odds the police have figured out that cemetery thing by now.\n",
        "\n",
        "(Enter EDITH)\n",
        "\n",
        "JEFF\n",
        "I hope so.\n",
        "\n",
        "EDITH\n",
        "If you're really that worried Jeff why don't you radio in and find out? Mac\n",
        "should be on duty at the field by now. He could call Paula and relay the message\n",
        "to you.\n",
        "\n",
        "DANNY\n",
        "Hi Edith.\n",
        "\n",
        "EDITH\n",
        "Hi Silents. I haven't heard a word from this end of the plane since we left the\n",
        "field.\n",
        "\n",
        "DANNY\n",
        "Jeff's been giving me and himself a study in silence.\n",
        "\n",
        "EDITH\n",
        "You boys are feudin'?\n",
        "\n",
        "JEFF\n",
        "Oh no Edie, nothing like that.\n",
        "\n",
        "DANNY\n",
        "Hey Edie, how about you and me balling it up in Albuquerque?\n",
        "\n",
        "EDITH\n",
        "Albuquerque? Have you read that flight schedule Boy?\n",
        "\n",
        "DANNY\n",
        "What about it?\n",
        "\n",
        "EDITH\n",
        "We land in Albuquerque at 4 am. That's strictly a nine o'clock town.\n",
        "\n",
        "DANNY\n",
        "Well I know a friend that'll help us --\n",
        "\n",
        "EDITH\n",
        "Let's have a problem first, huh Danny.\n",
        "\n",
        "DANNY\n",
        "Ah he's worried about Paula.\n",
        "\n",
        "EDITH\n",
        "I read about that cemetery business. I tried to get you kids to not buy too near\n",
        "one of those things. We get there soon enough as it is.\n",
        "\n",
        "DANNY\n",
        "He thought it'd be quiet and peaceful there.\n",
        "\n",
        "EDITH\n",
        "No doubt about that. It's quiet alright, like a tomb. I'm sorry Jeff, that was a\n",
        "bad joke.\n",
        "\n",
        "Using the example above and following description, write the dialogue of the scene.\n",
        "\n",
        "\"\"\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qiKa4EVKDbIZ"
      },
      "outputs": [],
      "source": [
        "#@title Custom\n",
        "\n",
        "#@markdown These prefixes for `CHARACTERS_PROMPT`, `SCENE_PROMPT`, `SETTING_PROMPT`, `TITLES_PROMPT` and `DIALOG_PROMPT` were written by the authors. They were not used in the evaluation study but can serve as a template to write custom prefix sets.\n",
        "\n",
        "#@markdown To write your own prompt prefix set, edit this code and pay attention to follow the existing formatting, with appropriate `STOP_MARKER`, `END_MARKER` and element markers.\n",
        "\n",
        "custom_prefixes = {}\n",
        "custom_prefixes['CHARACTERS_PROMPT'] = \"\"\"\n",
        "Here is an example of a logline and a list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"James finds a well in his backyard that is haunted by the ghost of Sam.\n",
        "\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\" James \"\"\" + DESCRIPTION_MARKER + \"\"\" James is twenty-six, serious about health and wellness and optimistic. \"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\" Sam \"\"\" + DESCRIPTION_MARKER + \"\"\" Sam fell down the well when he was 12, and was never heard from again. Sam is now a ghost. \"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Morgan adopts a new cat, Misterio, who sets a curse on anyone that pets them.\n",
        "\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\" Morgan \"\"\" + DESCRIPTION_MARKER + \"\"\" Morgan is booksmart and popular; they are trusting but also have been known to hold a grudge. \"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\" Misterio \"\"\" + DESCRIPTION_MARKER + \"\"\" Misterio is a beautiul black cat, it is of uncertain age; it has several gray whiskers that make it look wise and beyond its years.  \"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Mr. Dorbenson finds a book at a garage sale that tells the story of his own life. And it ends in a murder!\n",
        "\n",
        "\"\"\" + CHARACTER_MARKER + \"\"\" Mr. Glen Dorbenson \"\"\" + DESCRIPTION_MARKER + \"\"\" Mr. Glen Dorbenson frequents markets and garage sales always looking for a bargain. He is lonely and isolated and looking for his meaning in life. \"\"\" + STOP_MARKER + \"\"\"\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the examples above and the following logline, complete the list of characters.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "custom_prefixes['SCENE_PROMPT'] = \"\"\"\n",
        "Here is an example of a logline, a list of characters, and a list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"In the following story, James finds a well in his backyard that is haunted by the ghost of Sam. The main characters are James and Sam.\n",
        "James is twenty-six, serious about health and wellness and optimistic.\n",
        "Sam fell down the well when he was 12, and was never heard from again. Sam is now a ghost.\n",
        "\n",
        "\"\"\" + SCENES_MARKER + \"\"\"\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The backyard.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Beginning.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"James is weeding his garden in the backyard, the ghost of Sam is rummaging around in the well. James listens closely and hears the murmurs of Sam down the well. James unearths the opening to the well, and looks down to see a glimmering reflection.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The well.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Middle.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"James is making his way down the well, Sam's voice is reverberating on the walls of the well. Sam tells the story of how he came to haunt the well. James offers to help set the soul of Sam free.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The house.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Conclusion.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Looking at a photo of the gardden featuring Sam, James says his goodbyes to Sam, Sam thanks James for his help. The ghost of Sam is set free after and James goes living his life.\n",
        "\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Morgan adopts a new cat, Misterio, who sets a curse on anyone that pets them.\n",
        "The main characters are Morgan and Misterio (a cat).\n",
        "Morgan is booksmart and popular; they are trusting but also have been known to hold a grudge.\n",
        "Misterio is a beautiul black cat, it is of uncertain age; it has several gray whiskers that make it look wise and beyond its years.\n",
        "\n",
        "\"\"\" + SCENES_MARKER + \"\"\"\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The Adoption Center\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Beginning.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Morgan walks into The Adoption Center looking for a new pet. Morgan talks to the various cats and dogs in the center, they can hear a response from one very special cat: Misterio. Misterio is stuck in a cage. After sharing an interesting and intimate exchange, Morgan adopts Misterio on several conditions.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"Morgan's house.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Middle.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Morgan is describing to Misterio all the facts they know about felines, and then asks them to behave when company arrives. Misterio is getting pets from Morgan, broods and puurs with the pets of Morgan, they are up to something.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The back stoop.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Conclusion.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Morgan has gone to bed, and Misterio transtransmorgifies into a half-cat-half-human horror. Misterio wakes up Morgan with a meow loud enough to shatter the window. Morgan erupts from bed, realizing the consequences of their recent adoption and quickly try to fix things.\n",
        "\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and the following logline and list of characters, complete the list of plot points.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "custom_prefixes['SETTING_PROMPT'] = \"\"\"\n",
        "Here are examples of logline, location, and that location's description.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Morgan adopts a new cat, Misterio, who sets a curse on anyone that pets them.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The Adoption Center.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The Adoption Center is a sad place, especially for an unadopted pet. It is full of walls and walls of cages and cages. Inside of each is an abandoned animal, longing for a home. The lighting is dim, gray, buzzing fluorescent.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"James finds a well in his backyard that is haunted by the ghost of Sam.\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The well.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The well is buried under grass and hedges. It is at least twenty feet deep, if not more and it is masoned with stones. It is 150 years old at least. It stinks of stale, standing water, and has vines growing up the sides. It is narrow enough to not be able to fit down if you are a grown adult human.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_MARKER + \"\"\"Mr. Dorbenson finds a book at a garage sale that tells the story of his own life. And it ends in a murder!\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The garage sale.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"It is a garage packed with dusty household goods and antiques. There is a box at the back that says FREE and is full of paper back books.\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the examples above and the following logine and location name, complete location description.\n",
        "\n",
        "\"\"\" + LOGLINE_MARKER\n",
        "\n",
        "\n",
        "custom_prefixes['TITLES_PROMPT'] = \"\"\"\n",
        "Examples of alternative, original and descriptive titles for known play and film scripts.\n",
        "\n",
        "Example 1.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Bob has an argument with his best friend, Charles.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"The End of A Friend\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 2.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Terence tries and fails to become a wizard.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"Spellcaster\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 3.\n",
        "\"\"\" + LOGLINE_ELEMENT + \"\"\"Tom falls in love with Daisy.\n",
        "\"\"\" + TITLE_ELEMENT + \"\"\"The Greatest Love Story Ever Told\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Example 4.\n",
        "\"\"\" + LOGLINE_ELEMENT\n",
        "\n",
        "\n",
        "# Alternative summary, if concatenating stories and beats.\n",
        "# \"\"\" + SUMMARY_ELEMENT + \"\"\"\n",
        "\n",
        "custom_prefixes['DIALOG_PROMPT'] = \"\"\"\n",
        "Here is an example of description and scene dialogue from a modern screenplay.\n",
        "\n",
        "\"\"\" + PLACE_ELEMENT + \"\"\"The Adoption Center.\n",
        "\"\"\" + DESCRIPTION_ELEMENT + \"\"\"The Adoption Center is a sad place, especially for an unadopted pet. It is full of walls and walls of cages and cages. Inside of each is an abandoned animal, longing for a home. The lighting is dim, gray, buzzing fluorescent.\n",
        "\"\"\" + CHARACTERS_ELEMENT + \"\"\"Morgan is booksmart and popular; they are trusting but also have been known to hold a grudge.\n",
        "Misterio is a beautiul black cat, it is of uncertain age; it has several gray whiskers that make it look wise and beyond its years.\n",
        "\"\"\" + PLOT_ELEMENT + \"\"\"Beginning.\n",
        "\"\"\" + SUMMARY_ELEMENT + \"\"\"Morgan adopts a new cat, Misterio, who sets a curse on anyone that pets them.\n",
        "\"\"\" + BEAT_ELEMENT + \"\"\"Morgan walks into The Adoption Center looking for a new pet. Morgan talks to the various cats and dogs in the center, they can hear a response from one very special cat: Misterio. After sharing an interesting and intimate exchange, Morgan adopts Misterio on several conditions.\n",
        "\n",
        "\"\"\" + DIALOG_MARKER + \"\"\"\n",
        "\n",
        "MORGAN\n",
        "Well, well, well ... aren't you the most precious little rascal.\n",
        "\n",
        "Cats are meowing and dogs are barking. There is a loud purr in the background.\n",
        "\n",
        "MORGAN\n",
        "Look at this little face... how could you not love a little Devon Rex face like this. With whiskers almost as long as your tail.\n",
        "\n",
        "Morgan makes their way down the hallways, running their hand along the cages. They feel a warm fuzzy paw bat their fingers.\n",
        "\n",
        "MORGAN\n",
        "Hello precious, and what is your name?\n",
        "\n",
        "Misterio let's out a long and sustained meow.\n",
        "\n",
        "MORGAN\n",
        "Well, well, I am Morgan and it is nice to meet you.\n",
        "\n",
        "MISTERIO\n",
        "(meowing louder this time) purrr, purrr, purrr.\n",
        "\n",
        "Morgan reads the sign on the bottom right of the cage, it reads: Misterio.\n",
        "\n",
        "MORGAN\n",
        "You have the most amazing face, and beautiful eyes. I could absolutely get lost in them.\n",
        "\n",
        "Morgan and Misterio start to stare at each other. They look deeply into each others eyes. They start to breath in rhythm.\n",
        "\n",
        "MISTERIO\n",
        "I can hear what you are thinking...\n",
        "\n",
        "Morgan is startled and looks around to see if anyone else can hear the cat's thoughts...\n",
        "\n",
        "MORGAN\n",
        "(looking around) you can hear my thoughts?\n",
        "\n",
        "MISTERIO\n",
        "I can hear what you are thinking.\n",
        "\n",
        "MORGAN\n",
        "What?\n",
        "\n",
        "MISTERIO\n",
        "Yes, I can hear your thoughts.\n",
        "\n",
        "MORGAN\n",
        "You are amazing. Want to come home with me? Want your new forever home?\n",
        "\n",
        "MISTERIO\n",
        "Yes, I would love that.\n",
        "\n",
        "MISTERIO purrs loud enough that the other animals all fall silent.\n",
        "\n",
        "MORGAN\n",
        "I will adopt you on a few conditions. First, you must not talk to me at night when I am sleeping. Second, you must not talk to me when I am out in public.\n",
        "\n",
        "MISTERIO\n",
        "Okay.\n",
        "\n",
        "MORGAN\n",
        "Okay, it's a deal.\n",
        "\n",
        "Misterio runs around the cage, Morgan laughs as Misterio rubs against the cage and tries to jump in Morgan's arms as soon as the cage is opened.\n",
        "\"\"\" + END_MARKER + \"\"\"\n",
        "\n",
        "Using the example above and following description, write the dialogue of the scene.\n",
        "\n",
        "\"\"\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup (continued): configure the Large Language Model (LLM)"
      ],
      "metadata": {
        "id": "b1zuT88kLcqW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Google PaLM 2 and Gemini Language API\n",
        "\n",
        "#@markdown This cell contain code for using [Google PaLM 2](https://ai.google/discover/palm2/)\n",
        "#@markdown as language model for Dramatron.<br>\n",
        "#@markdown <br>\n",
        "#@markdown * **To use the Palm 2 API**<br>\n",
        "#@markdown You will first need to create a project called `PROJECT_ID` on the [Google Cloud Platform](https://cloud.google.com/) account,\n",
        "#@markdown then enable the [Vertex AI API](https://cloud.google.com/vertex-ai).<br>\n",
        "#@markdown To run this cell, you will need to authenticate by logging in to the same Google account as your `PROJECT_ID` GCP project.<br>\n",
        "#@markdown <br>\n",
        "#@markdown * **To use the Gemini API**<br>\n",
        "#@markdown Before you can use the Gemini API, you must first obtain an API key. If you don't already have one, create a key with one click in Google AI Studio.\n",
        "#@markdown <a class=\"button button-primary\" href=\"https://makersuite.google.com/app/apikey\" target=\"_blank\" rel=\"noopener noreferrer\">Get an API key</a>\n",
        "\n",
        "#@markdown See [Getting Started with the Vertex AI PaLM API & Python SDK](https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/language/intro_palm_api.ipynb)\n",
        "#@markdown for examples of how to call the API from Python in a colab,\n",
        "#@markdown and [Quickstart using the Vertex AI API](https://cloud.google.com/vertex-ai/docs/generative-ai/start/quickstarts/api-quickstart) for command-line requests from the Google Cloud Shell.\n",
        "#@markdown See [Gemini API: Quickstart with Python](https://colab.sandbox.google.com/github/google/generative-ai-docs/blob/main/site/en/tutorials/python_quickstart.ipynb) for examples of how to call the Gemini API.\n",
        "\n",
        "# Used by Gemini.\n",
        "import google.generativeai as genai\n",
        "\n",
        "# Used by Palm 2.\n",
        "import vertexai\n",
        "from vertexai.preview.language_models import TextGenerationModel\n",
        "\n",
        "# Initialise the Vertex AI with your Google credentials.\n",
        "GCP_PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "GOOGLE_API_KEY = \"\"  # @param {type:\"string\"}\n",
        "GOOGLE_MODEL_NAME = \"gemini-pro\" #@param [\"gemini-pro\", \"text-bison-32k\", \"text-bison@001\"]\n",
        "\n",
        "if len(GCP_PROJECT_ID) > 0:\n",
        "  from google.colab import auth\n",
        "  auth.authenticate_user()\n",
        "  vertexai.init(project=GCP_PROJECT_ID, location=\"us-central1\")\n",
        "if len(GOOGLE_API_KEY) > 0:\n",
        "  genai.configure(api_key=GOOGLE_API_KEY)\n",
        "GOOGLE_SAMPLING_TOP_K = 40\n",
        "\n",
        "\n",
        "class GoogleAPI(LanguageAPI):\n",
        "  \"\"\"A class wrapping the Google PaLM 2 language model API.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               sample_length: int,\n",
        "               model: Optional[str] = None,\n",
        "               model_param: Optional[str] = None,\n",
        "               config_sampling: Optional[dict] = None,\n",
        "               seed: Optional[int] = None,\n",
        "               max_retries: int = _MAX_RETRIES,\n",
        "               timeout: float = _TIMEOUT):\n",
        "    \"\"\"Initializer.\n",
        "\n",
        "    Args:\n",
        "      sample_length: Length of text to sample from model.\n",
        "      model: The model name to correct to. An error will be raised if it does\n",
        "        not exist.\n",
        "      model_param: Custom language model params.\n",
        "      config_sampling: ConfigDict with parameters.\n",
        "      seed: Random seed for sampling.\n",
        "      max_retries: Maximum number of retries for the remote API.\n",
        "      timeout: Maximum waiting timeout\n",
        "    \"\"\"\n",
        "    super().__init__(sample_length=sample_length,\n",
        "                     model=model,\n",
        "                     model_param=model_param,\n",
        "                     config_sampling=config_sampling,\n",
        "                     seed=seed,\n",
        "                     max_retries=max_retries,\n",
        "                     timeout=timeout)\n",
        "    if 'bison' in self._model:\n",
        "      self._client = TextGenerationModel.from_pretrained(model)\n",
        "    elif 'gemini' in self._model:\n",
        "      self._client = genai.GenerativeModel(model)\n",
        "\n",
        "  @property\n",
        "  def client(self):\n",
        "    return self._client\n",
        "\n",
        "  @property\n",
        "  def model_metadata(self):\n",
        "    return {'engine': self._model,\n",
        "            'model_param': self._model_param,\n",
        "            'max_tokens': self._sample_length}\n",
        "\n",
        "  def sample(self,\n",
        "             prompt: str,\n",
        "             sample_length: Optional[int] = None,\n",
        "             seed: Optional[int] = None,\n",
        "             num_samples: int = 1):\n",
        "    \"\"\"Sample model with provided prompt and optional sample_length and seed.\"\"\"\n",
        "    if sample_length is None:\n",
        "      sample_length = self._sample_length\n",
        "    if 'bison' in self._model:\n",
        "      response = self._client.predict(\n",
        "          prompt=prompt,\n",
        "          max_output_tokens=sample_length,\n",
        "          temperature=self._config_sampling['temp'],\n",
        "          top_p=self._config_sampling['prob'],\n",
        "          top_k=self._config_sampling['top_k'])\n",
        "    if 'gemini' in self._model:\n",
        "      response = self._client.generate_content(prompt)\n",
        "    results = [LanguageResponse(text=response.text,\n",
        "                                text_length=len(response.text),\n",
        "                                prompt=prompt,\n",
        "                                prompt_length=len(prompt))]\n",
        "    return results\n",
        "\n",
        "\n",
        "# Create the config.\n",
        "config = {}\n",
        "config['language_api_name'] = 'PaLM 2 / Gemini'\n",
        "config['model_param'] = GCP_PROJECT_ID\n",
        "config['model_name'] = GOOGLE_MODEL_NAME\n",
        "config['max_retries'] = MAX_RETRIES\n",
        "config['sample_length'] = SAMPLE_LENGTH\n",
        "config['max_paragraph_length'] = MAX_PARAGRAPH_LENGTH\n",
        "config['max_paragraph_length_characters'] = MAX_PARAGRAPH_LENGTH_CHARACTERS\n",
        "config['max_paragraph_length_scenes'] = MAX_PARAGRAPH_LENGTH_SCENES\n",
        "config['sampling'] = {}\n",
        "config['sampling']['prob'] = SAMPLING_PROB\n",
        "config['sampling']['temp'] = SAMPLING_TEMP\n",
        "config['sampling']['top_k'] = GOOGLE_SAMPLING_TOP_K\n",
        "config['prefixes'] = {}\n",
        "config['file_dir'] = None\n",
        "\n",
        "print('Config:')\n",
        "for key, value in config.items():\n",
        "  if key != 'prefixes':\n",
        "    print(f'{key}: {value}')\n",
        "\n",
        "client = GoogleAPI(\n",
        "    model_param=config['model_param'],\n",
        "    model=config['model_name'],\n",
        "    seed=DEFAULT_SEED,\n",
        "    sample_length=config['sample_length'],\n",
        "    max_retries=config['max_retries'],\n",
        "    config_sampling=config['sampling'])\n",
        "\n",
        "print(f'Client model metadata: {client.model_metadata}')\n",
        "\n",
        "prompt = 'Once upon a time, there was'\n",
        "results = client.sample(prompt)\n",
        "if len(results) > 0 and isinstance(results[0], LanguageResponse):\n",
        "  print(f'\\nPrompt: {prompt}\\nResponse: {results[0].text}')"
      ],
      "metadata": {
        "id": "1TeChk16H4K0",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title OpenAI ChatGPT API {\"run\": \"auto\"}\n",
        "\n",
        "#@markdown This cell contain code for using the [OpenAI Chat Completion API](https://platform.openai.com/docs/guides/text-generation/chat-completions-api)\n",
        "#@markdown as language model for Dramatron.<br>\n",
        "#@markdown <br>\n",
        "#@markdown Before you can use the OpenAI API, you must first obtain an API key. If you don't already have one,\n",
        "#@markdown <a class=\"button button-primary\" href=\"https://platform.openai.com/\" target=\"_blank\" rel=\"noopener noreferrer\">create an account</a>\n",
        "#@markdown then go to <a class=\"button button-primary\" href=\"https://platform.openai.com/api-keys\" target=\"_blank\" rel=\"noopener noreferrer\">API keys</a>\n",
        "#@markdown and click on `+ Create new secret key`.\n",
        "\n",
        "#@markdown See [Text Generation Models / Chat Completion API](https://platform.openai.com/docs/guides/text-generation)\n",
        "#@markdown for examples of how to call the API from Python in a colab.\n",
        "\n",
        "!pip install openai\n",
        "\n",
        "from openai import OpenAI\n",
        "import os\n",
        "\n",
        "# Initialise the OpenAI API with your OpenAI credentials.\n",
        "OPENAI_API_KEY = \"\"  # @param {type:\"string\"}\n",
        "CHATGPT_MODEL_NAME = \"gpt-4-1106-preview\" #@param [\"gpt-4-1106-preview\", \"gpt-3.5-turbo\"]\n",
        "CHATGPT_SYSTEM_PROMPT = \"You are a helpful playwright assistant.\" #@param {type:\"string\"}\n",
        "CHATGPT_FREQUENCY_PENALTY = 0.2\n",
        "CHATGPT_PRESENCE_PENALTY = 0.2\n",
        "os.environ['OPENAI_API_KEY'] = OPENAI_API_KEY\n",
        "\n",
        "\n",
        "class OpenAIAPI(LanguageAPI):\n",
        "  \"\"\"A class wrapping the OpenAI ChatGPT language model API.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               sample_length: int,\n",
        "               model: Optional[str] = None,\n",
        "               model_param: Optional[str] = None,\n",
        "               config_sampling: Optional[dict] = None,\n",
        "               seed: Optional[int] = None,\n",
        "               max_retries: int = _MAX_RETRIES,\n",
        "               timeout: float = _TIMEOUT):\n",
        "    \"\"\"Initializer.\n",
        "\n",
        "    Args:\n",
        "      sample_length: Length of text to sample from model.\n",
        "      model: The model name to correct to. An error will be raised if it does\n",
        "        not exist.\n",
        "      model_param: Custom language model params.\n",
        "      config_sampling: ConfigDict with parameters.\n",
        "      seed: Random seed for sampling.\n",
        "      max_retries: Maximum number of retries for the remote API.\n",
        "      timeout: Maximum waiting timeout\n",
        "    \"\"\"\n",
        "    super().__init__(sample_length=sample_length,\n",
        "                     model=model,\n",
        "                     model_param=model_param,\n",
        "                     config_sampling=config_sampling,\n",
        "                     seed=seed,\n",
        "                     max_retries=max_retries,\n",
        "                     timeout=timeout)\n",
        "    self._client = OpenAI()\n",
        "\n",
        "  @property\n",
        "  def client(self):\n",
        "    return self._client\n",
        "\n",
        "  @property\n",
        "  def model_metadata(self):\n",
        "    return {'engine': self._model,\n",
        "            'model_param': self._model_param,\n",
        "            'max_tokens': self._sample_length}\n",
        "\n",
        "  def sample(self,\n",
        "             prompt: str,\n",
        "             sample_length: Optional[int] = None,\n",
        "             seed: Optional[int] = None,\n",
        "             num_samples: int = 1):\n",
        "    \"\"\"Sample model with provided prompt and optional sample_length and seed.\"\"\"\n",
        "    if sample_length is None:\n",
        "      sample_length = self._sample_length\n",
        "    response = self._client.chat.completions.create(\n",
        "        model=self._model,\n",
        "        max_tokens=sample_length,\n",
        "        temperature=self._config_sampling['temp'],\n",
        "        top_p=self._config_sampling['prob'],\n",
        "        frequency_penalty=self._config_sampling['frequency_penalty'],\n",
        "        presence_penalty=self._config_sampling['presence_penalty'],\n",
        "        messages=[\n",
        "          {\"role\": \"system\", \"content\": self._model_param},\n",
        "          {\"role\": \"user\", \"content\": prompt}\n",
        "        ]\n",
        "    )\n",
        "    response_text = ''\n",
        "    if len(response.choices) > 0:\n",
        "      response_text = response.choices[0].message.content\n",
        "    results = [LanguageResponse(text=response_text,\n",
        "                                text_length=len(response_text),\n",
        "                                prompt=prompt,\n",
        "                                prompt_length=len(prompt))]\n",
        "    return results\n",
        "\n",
        "\n",
        "# Create the config.\n",
        "config = {}\n",
        "config['language_api_name'] = 'OpenAI'\n",
        "config['model_param'] = CHATGPT_SYSTEM_PROMPT\n",
        "config['model_name'] = CHATGPT_MODEL_NAME\n",
        "config['max_retries'] = MAX_RETRIES\n",
        "config['sample_length'] = SAMPLE_LENGTH\n",
        "config['max_paragraph_length'] = MAX_PARAGRAPH_LENGTH\n",
        "config['max_paragraph_length_characters'] = MAX_PARAGRAPH_LENGTH_CHARACTERS\n",
        "config['max_paragraph_length_scenes'] = MAX_PARAGRAPH_LENGTH_SCENES\n",
        "config['sampling'] = {}\n",
        "config['sampling']['prob'] = SAMPLING_PROB\n",
        "config['sampling']['temp'] = SAMPLING_TEMP\n",
        "config['sampling']['frequency_penalty'] = CHATGPT_FREQUENCY_PENALTY\n",
        "config['sampling']['presence_penalty'] = CHATGPT_PRESENCE_PENALTY\n",
        "config['prefixes'] = {}\n",
        "config['file_dir'] = None\n",
        "\n",
        "print('Config:')\n",
        "for key, value in config.items():\n",
        "  if key != 'prefixes':\n",
        "    print(f'{key}: {value}')\n",
        "\n",
        "client = OpenAIAPI(\n",
        "    model_param=config['model_param'],\n",
        "    model=config['model_name'],\n",
        "    seed=DEFAULT_SEED,\n",
        "    sample_length=config['sample_length'],\n",
        "    max_retries=config['max_retries'],\n",
        "    config_sampling=config['sampling'])\n",
        "\n",
        "print(f'Client model metadata: {client.model_metadata}')\n",
        "\n",
        "prompt = 'Once upon a time, there was'\n",
        "results = client.sample(prompt, sample_length=256)\n",
        "if len(results) > 0 and isinstance(results[0], LanguageResponse):\n",
        "  print(f'\\nPrompt: {prompt}\\nResponse: {results[0].text}')"
      ],
      "metadata": {
        "id": "rNhsuyTR6AoS",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Mixtral API via Groq {\"run\": \"auto\"}\n",
        "\n",
        "#@markdown This cell contain code for using the [Groq API for chat completion](https://platform.openai.com/docs/guides/text-generation/chat-completions-api)\n",
        "#@markdown for Dramatron, with several possible language models including [Mixtral 8x7B](https://mistral.ai/news/mixtral-of-experts/).<br>\n",
        "#@markdown <br>\n",
        "#@markdown Before you can use the Groq API, you must first obtain an API key. If you don't already have one,\n",
        "#@markdown <a class=\"button button-primary\" href=\"https://console.groq.com/\" target=\"_blank\" rel=\"noopener noreferrer\">create an account</a>\n",
        "#@markdown then go to <a class=\"button button-primary\" href=\"https://console.groq.com/keys\" target=\"_blank\" rel=\"noopener noreferrer\">API keys</a>\n",
        "#@markdown and click on `Create API key`.\n",
        "\n",
        "#@markdown See [GroqCloud Documentation](https://console.groq.com/docs/quickstart)\n",
        "#@markdown for examples of how to call the API from Python in a colab.\n",
        "\n",
        "!pip install groq\n",
        "\n",
        "from groq import Groq\n",
        "import os\n",
        "\n",
        "# Initialise the Groq API with your OpenAI credentials.\n",
        "GROQ_API_KEY = \"\"  # @param {type:\"string\"}\n",
        "GROQ_MODEL_NAME = \"mixtral-8x7b-32768\" #@param [\"mixtral-8x7b-32768\"]\n",
        "GROQ_SYSTEM_PROMPT = \"You are a creative writing assistant for a team of writers. Your goal is to expand on the input text prompt and to generate the continuation of that text without any comments. Be as creative as possible, write rich detailed descriptions and use precise language. Add new original ideas. Finish generation with **END**.\" #@param {type:\"string\"}\n",
        "GROQ_FREQUENCY_PENALTY = 0.2\n",
        "GROQ_PRESENCE_PENALTY = 0.2\n",
        "os.environ['GROQ_API_KEY'] = GROQ_API_KEY\n",
        "\n",
        "\n",
        "class GroqAPI(LanguageAPI):\n",
        "  \"\"\"A class wrapping the Groq language model API.\"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               sample_length: int,\n",
        "               model: Optional[str] = None,\n",
        "               model_param: Optional[str] = None,\n",
        "               config_sampling: Optional[dict] = None,\n",
        "               seed: Optional[int] = None,\n",
        "               max_retries: int = _MAX_RETRIES,\n",
        "               timeout: float = _TIMEOUT):\n",
        "    \"\"\"Initializer.\n",
        "\n",
        "    Args:\n",
        "      sample_length: Length of text to sample from model.\n",
        "      model: The model name to correct to. An error will be raised if it does\n",
        "        not exist.\n",
        "      model_param: Custom language model params.\n",
        "      config_sampling: ConfigDict with parameters.\n",
        "      seed: Random seed for sampling.\n",
        "      max_retries: Maximum number of retries for the remote API.\n",
        "      timeout: Maximum waiting timeout\n",
        "    \"\"\"\n",
        "    super().__init__(sample_length=sample_length,\n",
        "                     model=model,\n",
        "                     model_param=model_param,\n",
        "                     config_sampling=config_sampling,\n",
        "                     seed=seed,\n",
        "                     max_retries=max_retries,\n",
        "                     timeout=timeout)\n",
        "    self._client = Groq(api_key=os.environ.get(\"GROQ_API_KEY\"))\n",
        "\n",
        "  @property\n",
        "  def client(self):\n",
        "    return self._client\n",
        "\n",
        "  @property\n",
        "  def model_metadata(self):\n",
        "    return {'engine': self._model,\n",
        "            'model_param': self._model_param,\n",
        "            'max_tokens': self._sample_length}\n",
        "\n",
        "  def sample(self,\n",
        "             prompt: str,\n",
        "             sample_length: Optional[int] = None,\n",
        "             seed: Optional[int] = None,\n",
        "             num_samples: int = 1):\n",
        "    \"\"\"Sample model with provided prompt and optional sample_length and seed.\"\"\"\n",
        "    if sample_length is None:\n",
        "      sample_length = self._sample_length\n",
        "    response = self._client.chat.completions.create(\n",
        "        model=self._model,\n",
        "        max_tokens=sample_length,\n",
        "        temperature=self._config_sampling['temp'],\n",
        "        top_p=self._config_sampling['prob'],\n",
        "        messages=[\n",
        "          {\"role\": \"system\", \"content\": self._model_param},\n",
        "          {\"role\": \"user\", \"content\": prompt}\n",
        "        ]\n",
        "    )\n",
        "    response_text = ''\n",
        "    if len(response.choices) > 0:\n",
        "      response_text = response.choices[0].message.content\n",
        "    results = [LanguageResponse(text=response_text,\n",
        "                                text_length=len(response_text),\n",
        "                                prompt=prompt,\n",
        "                                prompt_length=len(prompt))]\n",
        "    return results\n",
        "\n",
        "\n",
        "# Create the config.\n",
        "config = {}\n",
        "config['language_api_name'] = 'Groq'\n",
        "config['model_param'] = GROQ_SYSTEM_PROMPT\n",
        "config['model_name'] = GROQ_MODEL_NAME\n",
        "config['max_retries'] = MAX_RETRIES\n",
        "config['sample_length'] = SAMPLE_LENGTH\n",
        "config['max_paragraph_length'] = MAX_PARAGRAPH_LENGTH\n",
        "config['max_paragraph_length_characters'] = MAX_PARAGRAPH_LENGTH_CHARACTERS\n",
        "config['max_paragraph_length_scenes'] = MAX_PARAGRAPH_LENGTH_SCENES\n",
        "config['sampling'] = {}\n",
        "config['sampling']['prob'] = SAMPLING_PROB\n",
        "config['sampling']['temp'] = SAMPLING_TEMP\n",
        "config['prefixes'] = {}\n",
        "config['file_dir'] = None\n",
        "\n",
        "print('Config:')\n",
        "for key, value in config.items():\n",
        "  if key != 'prefixes':\n",
        "    print(f'{key}: {value}')\n",
        "\n",
        "client = GroqAPI(\n",
        "    model_param=config['model_param'],\n",
        "    model=config['model_name'],\n",
        "    seed=DEFAULT_SEED,\n",
        "    sample_length=config['sample_length'],\n",
        "    max_retries=config['max_retries'],\n",
        "    config_sampling=config['sampling'])\n",
        "\n",
        "print(f'Client model metadata: {client.model_metadata}')\n",
        "\n",
        "prompt = 'Once upon a time, there was'\n",
        "results = client.sample(prompt, sample_length=256)\n",
        "if len(results) > 0 and isinstance(results[0], LanguageResponse):\n",
        "  print(f'\\nPrompt: {prompt}\\nResponse: {results[0].text}')"
      ],
      "metadata": {
        "id": "xU5KZgFIrRb6",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mVvjeEzSRoXq"
      },
      "source": [
        "# Choose a prompt prefix set"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JsYTDv-5CnCF"
      },
      "outputs": [],
      "source": [
        "#@title Choose among prefix sets { run: \"auto\" }\n",
        "\n",
        "prefix_set = 'medea_prefixes' #@param ['medea_prefixes', 'scifi_prefixes', 'custom_prefixes']\n",
        "prefixes = eval(prefix_set)\n",
        "config['prefixes'] = prefixes\n",
        "\n",
        "print(f'Loaded {prefix_set}.')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7_vuAy43H6J-"
      },
      "source": [
        "# Interactive story generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OfnN_7DqhQx7"
      },
      "outputs": [],
      "source": [
        "#@title Define the **Log line** { run: \"auto\" }\n",
        "\n",
        "#@markdown Log lines are one- or two-sentence summaries of the action.\n",
        "#@markdown They typically contain the **setting**, **protagonist**, **antagonist**, a **conflict** or **goal** and sometimes the **inciting incident**.\n",
        "logline = \"Folk tale about a rabbit, a fox and a crow living in an enchanted forest. The cunning animals safeguard the golden apple tree from a greedy lumberjack and conspire to hide the lumberjack's axe.\" #@param {type:\"string\"}\n",
        "print(logline)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "renVLifoRbhc"
      },
      "outputs": [],
      "source": [
        "#@title Create the story generator\n",
        "\n",
        "#@markdown WARNING: running this cell resets the generator and all its outputs.\n",
        "\n",
        "generator = StoryGenerator(\n",
        "    storyline=logline,\n",
        "    prefixes=prefixes,\n",
        "    max_paragraph_length=config['max_paragraph_length'],\n",
        "    client=client,\n",
        "    filter=filter)\n",
        "\n",
        "print(f'New Dramatron generator created.')\n",
        "\n",
        "story = None"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oRaJlEme-phR"
      },
      "source": [
        "## Usage\n",
        "Running each cell displays the UI and generates the first suggestion.\n",
        "\n",
        "Once you have run the cell once, you can do the following actions _without re-running the cell_:\n",
        "* Click on `Generate new` to **generate a new suggestion**.\n",
        "* Once you generated a suggestion, you can **edit the suggestion** in the text box. It is automatically saved and used in the next step.\n",
        "* Click `Continue generating` to **add to the suggestion**.\n",
        "* Once you have generated multiple suggestions, you can **navigate through the suggestion history** by clicking on `Previous` and `Next`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3UYrzDgbspAG"
      },
      "outputs": [],
      "source": [
        "#@title Generate a **Title**\n",
        "\n",
        "data_title = {\"text\": \"\", \"text_area\": None, \"seed\": generator.seed - 1}\n",
        "\n",
        "def fun_generate_title(button):\n",
        "  data_title[\"seed\"] += 1\n",
        "  seed = data_title[\"seed\"]\n",
        "  new_title_button.description = f\"Generating {seed}...\"\n",
        "  generator.step(0, seed=seed)\n",
        "  data_title[\"text\"] = generator.title_str().strip()\n",
        "  new_title_button.description = \"Generate new\"\n",
        "  if data_title[\"text_area\"] is not None:\n",
        "    data_title[\"text_area\"].value = data_title[\"text\"]\n",
        "\n",
        "def fun_rewrite_title(text):\n",
        "  text_to_parse = TITLE_ELEMENT + text + END_MARKER\n",
        "  generator.rewrite(text_to_parse, level=1)\n",
        "  return text\n",
        "\n",
        "# Widget to generate new title.\n",
        "new_title_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Generate new', tooltip='Generate new', disabled=False)\n",
        "new_title_button.on_click(fun_generate_title)\n",
        "display(new_title_button)\n",
        "\n",
        "# Widget to rewrite the title.\n",
        "layout = widgets.Layout(height='50px', min_height='60px', width='auto')\n",
        "data_title[\"text_area\"] = widgets.Textarea(\n",
        "    value=data_title[\"text\"], layout=layout, description=' ',\n",
        "    style={'description_width': 'initial'})\n",
        "textarea_title = widgets.interactive(\n",
        "    fun_rewrite_title, text=data_title[\"text_area\"])\n",
        "display(textarea_title)\n",
        "data_title[\"text_area\"].value = data_title[\"text\"]\n",
        "fun_generate_title(new_title_button)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "yq7y3NIUYtD8"
      },
      "outputs": [],
      "source": [
        "#@title Generate **Characters**\n",
        "\n",
        "data_chars = {\"text\": \"\", \"text_area\": None, \"seed\": generator.seed - 1,\n",
        "              \"history\": GenerationHistory(), \"lock\": False}\n",
        "\n",
        "def fun_generate_characters(_):\n",
        "  data_chars[\"seed\"] += 1\n",
        "  seed = data_chars[\"seed\"]\n",
        "  data_chars[\"lock\"] = True\n",
        "  while True:\n",
        "    new_characters_button.description = f\"Generating {seed}...\"\n",
        "    generator.step(1, seed=seed)\n",
        "    data_chars[\"text\"] = strip_remove_end(generator.characters.to_string())\n",
        "    # Test if characters were actually generated.\n",
        "    if len(data_chars[\"text\"]) == 0:\n",
        "      seed += 1\n",
        "    else:\n",
        "      break\n",
        "  data_chars[\"seed\"] = seed\n",
        "  data_chars[\"history\"].add(data_chars[\"text\"], GenerationAction.NEW)\n",
        "  new_characters_button.description = \"Generate new\"\n",
        "  if data_chars[\"text_area\"] is not None:\n",
        "    data_chars[\"text_area\"].value = data_chars[\"text\"]\n",
        "  data_chars[\"lock\"] = False\n",
        "\n",
        "def fun_continue_characters(_):\n",
        "  data_chars[\"seed\"] += 1\n",
        "  seed = data_chars[\"seed\"]\n",
        "  data_chars[\"lock\"] = True\n",
        "  characters_continue_button.description = f\"Generating {seed}...\"\n",
        "  generator.complete(level=2, seed=seed, sample_length=256)\n",
        "  data_chars[\"text\"] = strip_remove_end(generator.characters.to_string())\n",
        "  data_chars[\"history\"].add(data_chars[\"text\"], GenerationAction.CONTINUE)\n",
        "  characters_continue_button.description = \"Continue generation\"\n",
        "  if data_chars[\"text_area\"] is not None:\n",
        "    data_chars[\"text_area\"].value = data_chars[\"text\"]\n",
        "  data_chars[\"lock\"] = False\n",
        "\n",
        "def fun_back_forward(data, history: GenerationHistory, delta: int):\n",
        "  data[\"lock\"] = True\n",
        "  if delta > 0:\n",
        "    data[\"text\"] = history.next()\n",
        "  if delta < 0:\n",
        "    data[\"text\"] = history.previous()\n",
        "  if data[\"text\"] is not None and data[\"text_area\"] is not None:\n",
        "      data[\"text_area\"].value = data[\"text\"]\n",
        "  data[\"lock\"] = False\n",
        "\n",
        "def fun_back_characters(_):\n",
        "  fun_back_forward(data_chars, data_chars[\"history\"], -1)\n",
        "\n",
        "def fun_forward_characters(_):\n",
        "  fun_back_forward(data_chars, data_chars[\"history\"], 1)\n",
        "\n",
        "def fun_rewrite_characters(text):\n",
        "  data_chars[\"text\"] = text\n",
        "  text_to_parse = text + END_MARKER\n",
        "  generator.rewrite(text_to_parse, level=2)\n",
        "  if data_chars[\"lock\"] is False:\n",
        "    data_chars[\"history\"].add(text, GenerationAction.REWRITE)\n",
        "  return text\n",
        "\n",
        "# Widget to generate new characters.\n",
        "new_characters_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Generate new', tooltip='Generate new', disabled=False)\n",
        "new_characters_button.on_click(fun_generate_characters)\n",
        "# Widget to continue the generation of the current characters.\n",
        "characters_continue_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Continue generation', tooltip='Continue generation',\n",
        "    disabled=False)\n",
        "characters_continue_button.on_click(fun_continue_characters)\n",
        "# Widgets to move back and forward in history of generation.\n",
        "back_characters_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Previous', tooltip='Previous', disabled=False)\n",
        "back_characters_button.on_click(fun_back_characters)\n",
        "forward_characters_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Next', tooltip='Next', disabled=False)\n",
        "forward_characters_button.on_click(fun_forward_characters)\n",
        "# Organise the widgets.\n",
        "display(widgets.HBox([new_characters_button, characters_continue_button,\n",
        "                      back_characters_button, forward_characters_button]))\n",
        "\n",
        "# Render the characters using widgets.\n",
        "layout = widgets.Layout(height='390px', min_height='400px', width='auto')\n",
        "data_chars[\"text_area\"] = widgets.Textarea(\n",
        "    value=data_chars[\"text\"], layout=layout, description=' ',\n",
        "    style={'description_width': 'initial'})\n",
        "textarea_chars = widgets.interactive(\n",
        "    fun_rewrite_characters, text=data_chars[\"text_area\"])\n",
        "display(textarea_chars)\n",
        "\n",
        "# Trigger generation for first seed.\n",
        "fun_generate_characters(new_characters_button)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KxtX6CxMb6tA"
      },
      "outputs": [],
      "source": [
        "#@title Generate a **Plot Synopsis** (sequence of **Scenes**)\n",
        "\n",
        "data_scenes = {\"text\": None, \"text_area\": None, \"seed\": generator.seed - 1,\n",
        "               \"history\": GenerationHistory(), \"lock\": False}\n",
        "\n",
        "def fun_generate_scenes(_):\n",
        "  data_scenes[\"seed\"] += 1\n",
        "  seed = data_scenes[\"seed\"]\n",
        "  data_scenes[\"lock\"] = True\n",
        "  new_scenes_button.description = f\"Generating {seed}...\"\n",
        "  generator.step(2, seed=seed)\n",
        "  data_scenes[\"text\"] = strip_remove_end(generator.scenes.to_string())\n",
        "  data_scenes[\"history\"].add(data_scenes[\"text\"], GenerationAction.NEW)\n",
        "  new_scenes_button.description = \"Generate new\"\n",
        "  if data_scenes[\"text_area\"] is not None:\n",
        "    data_scenes[\"text_area\"].value = data_scenes[\"text\"]\n",
        "  data_scenes[\"lock\"] = False\n",
        "\n",
        "def fun_continue_scenes(_):\n",
        "  data_scenes[\"seed\"] += 1\n",
        "  seed = data_scenes[\"seed\"]\n",
        "  data_scenes[\"lock\"] = True\n",
        "  scenes_continue_button.description = f\"Generating {seed}...\"\n",
        "  generator.complete(level=3, seed=seed, sample_length=256)\n",
        "  data_scenes[\"text\"] = strip_remove_end(generator.scenes.to_string())\n",
        "  data_scenes[\"history\"].add(data_scenes[\"text\"], GenerationAction.CONTINUE)\n",
        "  scenes_continue_button.description = \"Continue generation\"\n",
        "  if data_scenes[\"text_area\"] is not None:\n",
        "    data_scenes[\"text_area\"].value = data_scenes[\"text\"]\n",
        "  data_scenes[\"lock\"] = False\n",
        "\n",
        "def fun_back_scenes(_):\n",
        "  fun_back_forward(data_scenes, data_scenes[\"history\"], -1)\n",
        "\n",
        "def fun_forward_scenes(_):\n",
        "  fun_back_forward(data_scenes, data_scenes[\"history\"], 1)\n",
        "\n",
        "def fun_rewrite_scenes(text):\n",
        "  generator.rewrite(text, level=3)\n",
        "  if data_scenes[\"lock\"] is False:\n",
        "    data_scenes[\"history\"].add(text, GenerationAction.REWRITE)\n",
        "  return text\n",
        "\n",
        "# Widget to generate new scenes.\n",
        "new_scenes_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Generate new', tooltip='Generate new', disabled=False)\n",
        "new_scenes_button.on_click(fun_generate_scenes)\n",
        "# Widget to continue the generation of the current scenes.\n",
        "scenes_continue_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Continue generation', tooltip='Continue generation',\n",
        "    disabled=False)\n",
        "scenes_continue_button.on_click(fun_continue_scenes)\n",
        "# Widgets to move back and forward in history of generation.\n",
        "back_scenes_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Previous', tooltip='Previous', disabled=False)\n",
        "back_scenes_button.on_click(fun_back_scenes)\n",
        "forward_scenes_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Next', tooltip='Next', disabled=False)\n",
        "forward_scenes_button.on_click(fun_forward_scenes)\n",
        "# Organise the widgets.\n",
        "display(widgets.HBox([new_scenes_button, scenes_continue_button,\n",
        "                      back_scenes_button, forward_scenes_button]))\n",
        "\n",
        "# Render the scenes using widgets.\n",
        "layout = widgets.Layout(height='590px', min_height='600px', width='auto')\n",
        "data_scenes[\"text_area\"] = widgets.Textarea(\n",
        "    value=data_scenes[\"text\"], layout=layout, description=' ',\n",
        "    style={'description_width': 'initial'})\n",
        "scanes_textarea = widgets.interactive(\n",
        "    fun_rewrite_scenes, text=data_scenes[\"text_area\"])\n",
        "display(scanes_textarea)\n",
        "\n",
        "# Trigger generation for first seed.\n",
        "fun_generate_scenes(new_scenes_button)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YoBiBkIxQIRl"
      },
      "outputs": [],
      "source": [
        "#@title Generate **Place Descriptions**\n",
        "\n",
        "#@markdown This cell generates a description for each **Place** name in the **Plot Synopsis**. If you edit place names in the **Plot Synopsis** _after_ having run this cell, you will need to re-run this cell to update **Places**.\n",
        "\n",
        "place_names = list(set([scene.place for scene in generator.scenes[0]]))\n",
        "place_descriptions = {place_name: Place(place_name, '')\n",
        "                      for place_name in place_names}\n",
        "data_places = {\"descriptions\": place_descriptions, \"text_area\": {},\n",
        "               \"seed\": generator.seed - 1}\n",
        "\n",
        "def fun_generate_places(_):\n",
        "  data_places[\"seed\"] += 1\n",
        "  seed = data_places[\"seed\"]\n",
        "  # Update existing text areas with a waiting message.\n",
        "  new_places_button.description = f\"Generating {seed}...\"\n",
        "  # Generate all the places.\n",
        "  generator.step(3, seed=seed)\n",
        "  data_places[\"descriptions\"] = generator.places\n",
        "  new_places_button.description = \"Generate new\"\n",
        "  missing_places = {k: True for k in data_places[\"text_area\"].keys()}\n",
        "  for place_name, place_description in data_places[\"descriptions\"].items():\n",
        "    if place_name in data_places[\"text_area\"]:\n",
        "      description = place_description.description\n",
        "      data_places[\"text_area\"][place_name].value = description\n",
        "      del missing_places[place_name]\n",
        "    else:\n",
        "      print(f\"\\nWarning: [{place_name}] was added to the plot synopsis.\")\n",
        "      print(f\"Make a copy of the outputs and re-run the cell.\")\n",
        "  for place_name in missing_places:\n",
        "    data_places[\"text_area\"][place_name].value = (\n",
        "        f\"Warning: [{place_name}] was removed from the plot synopsis. \"\n",
        "        \"Make a copy of the outputs and re-run the cell.\")\n",
        "\n",
        "def fun_rewrite_places(place_name, text):\n",
        "  generator.rewrite(text, level=4, entity=place_name)\n",
        "  return text\n",
        "\n",
        "# Widget to generate new scenes.\n",
        "new_places_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Generate new', tooltip='Generate new', disabled=False)\n",
        "new_places_button.on_click(fun_generate_places)\n",
        "display(new_places_button)\n",
        "\n",
        "# Render each place using widgets.\n",
        "for place_name, place_description in data_places[\"descriptions\"].items():\n",
        "  text_place = place_description.description\n",
        "  layout = widgets.Layout(height='90px', min_height='100px', width='auto')\n",
        "  display(widgets.Label(place_name))\n",
        "  data_places[\"text_area\"][place_name] = widgets.Textarea(\n",
        "      value=text_place, layout=layout, description=' ',\n",
        "      style={'description_width': 'initial'})\n",
        "  textarea_place = widgets.interactive(\n",
        "      fun_rewrite_places, place_name=widgets.fixed(place_name),\n",
        "      text=data_places[\"text_area\"][place_name])\n",
        "  display(textarea_place)\n",
        "\n",
        "# Trigger generation for first seed.\n",
        "fun_generate_places(new_places_button)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "R36-qkMAhZ7a"
      },
      "outputs": [],
      "source": [
        "#@title Generate **Dialogues**\n",
        "\n",
        "num_scenes =  generator.num_scenes()\n",
        "\n",
        "data_dialogs = {\n",
        "    \"lock\": False, \n",
        "    \"text_area\": None, \n",
        "    \"seed\": generator.seed - 1,\n",
        "    \"history\": [GenerationHistory() for _ in range(99)], \n",
        "    \"scene\": 1\n",
        "}\n",
        "\n",
        "# Ensure idx_dialog is initialize\n",
        "idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "\n",
        "def fun_generate_dialog(_):\n",
        "  data_dialogs[\"seed\"] += 1\n",
        "  seed = data_dialogs[\"seed\"]\n",
        "  idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "  data_dialogs[\"lock\"] = True\n",
        "  new_dialog_button.description = f\"Generating {seed}...\"\n",
        "  generator.step(4, seed=seed, idx=idx_dialog)\n",
        "  data_dialogs[\"history\"][idx_dialog].add(\n",
        "      generator.dialogs[idx_dialog], GenerationAction.NEW)\n",
        "  new_dialog_button.description = \"Generate new\"\n",
        "  if data_dialogs[\"text_area\"] is not None:\n",
        "    data_dialogs[\"text_area\"].value = generator.dialogs[idx_dialog]\n",
        "  data_dialogs[\"lock\"] = False\n",
        "\n",
        "def fun_load_dialog(scene):\n",
        "  idx_dialog = scene - 1\n",
        "  scene_exists = (\n",
        "      len(generator.dialogs) > idx_dialog and\n",
        "      len(generator.dialogs[idx_dialog]) > 0)\n",
        "  # Update existing text area with a waiting message or load existing scene.\n",
        "  if scene_exists:\n",
        "    data_dialogs[\"lock\"] = True\n",
        "    if data_dialogs[\"text_area\"] is not None:\n",
        "      data_dialogs[\"text_area\"].value = generator.dialogs[idx_dialog]\n",
        "    data_dialogs[\"scene\"] = scene\n",
        "    data_dialogs[\"lock\"] = False\n",
        "  else:\n",
        "    data_dialogs[\"scene\"] = scene\n",
        "    fun_generate_dialog(None)\n",
        "\n",
        "def fun_continue_dialog(_):\n",
        "  data_dialogs[\"seed\"] += 1\n",
        "  seed = data_dialogs[\"seed\"]\n",
        "  idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "  data_dialogs[\"lock\"] = True\n",
        "  dialog_continue_button.description = f\"Generating {seed}...\"\n",
        "  generator.complete(level=5, seed=seed, entity=idx_dialog,\n",
        "                     sample_length=SAMPLE_LENGTH)\n",
        "  data_dialogs[\"history\"][idx_dialog].add(\n",
        "      generator.dialogs[idx_dialog], GenerationAction.CONTINUE)\n",
        "  dialog_continue_button.description = \"Continue generation\"\n",
        "  if data_dialogs[\"text_area\"] is not None:\n",
        "    data_dialogs[\"text_area\"].value = generator.dialogs[idx_dialog]\n",
        "  data_dialogs[\"lock\"] = False\n",
        "\n",
        "def fun_back_dialog(_):\n",
        "  idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "  if idx_dialog >= 0 and idx_dialog < len(data_dialogs[\"history\"]):\n",
        "    fun_back_forward(data_dialogs, data_dialogs[\"history\"][idx_dialog], -1)\n",
        "\n",
        "def fun_forward_dialog(_):\n",
        "  idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "  if idx_dialog >= 0 and idx_dialog < len(data_dialogs[\"history\"]):\n",
        "    fun_back_forward(data_dialogs, data_dialogs[\"history\"][idx_dialog], 1)\n",
        "\n",
        "# Function to edit the specific dialog.\n",
        "def fun_rewrite_dialog(text):\n",
        "  if data_dialogs[\"lock\"] == False:\n",
        "    idx_dialog = data_dialogs[\"scene\"] - 1\n",
        "    generator.rewrite(text, level=5, entity=idx_dialog)\n",
        "  return text\n",
        "\n",
        "# Widget to choose a seed and generate new scenes.\n",
        "scene_slider = widgets.IntSlider(\n",
        "    description='scene', min=1, max=generator.num_scenes())\n",
        "widgets.interactive(fun_load_dialog, scene=scene_slider)\n",
        "\n",
        "# Widget to generate new dialogue.\n",
        "new_dialog_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Generate new', tooltip='Generate new', disabled=False)\n",
        "new_dialog_button.on_click(fun_generate_dialog)\n",
        "# Widget to continue the generation of the current dialogue.\n",
        "dialog_continue_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Continue generation', tooltip='Continue generation',\n",
        "    disabled=False)\n",
        "dialog_continue_button.on_click(fun_continue_dialog)\n",
        "# Widgets to move back and forward in history of generation.\n",
        "back_dialog_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Previous', tooltip='Previous', disabled=False)\n",
        "back_dialog_button.on_click(fun_back_dialog)\n",
        "forward_dialog_button = widgets.Button(button_style='', icon='check',\n",
        "    description='Next', tooltip='Next', disabled=False)\n",
        "forward_dialog_button.on_click(fun_forward_dialog)\n",
        "# Organise the widgets.\n",
        "display(widgets.HBox([scene_slider, new_dialog_button, dialog_continue_button,\n",
        "                      back_dialog_button, forward_dialog_button]))\n",
        "\n",
        "# Render the dialog using widgets.\n",
        "layout = widgets.Layout(height='600px', width='auto')\n",
        "data_dialogs[\"text_area\"] = widgets.Textarea(\n",
        "    value=strip_remove_end(generator.dialogs[idx_dialog]), layout=layout,\n",
        "    description=' ')\n",
        "textarea_dialogs = widgets.interactive(\n",
        "    fun_rewrite_dialog, text=data_dialogs[\"text_area\"])\n",
        "display(textarea_dialogs)\n",
        "\n",
        "# Trigger generation for first seed.\n",
        "fun_generate_dialog(new_dialog_button)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-2MpykfCqO85"
      },
      "outputs": [],
      "source": [
        "#@title Render the script\n",
        "\n",
        "#@markdown Run this cell to render the whole story as a text string. This cells also renders all the edits and prefixes as text strings.\n",
        "\n",
        "# Render the story.\n",
        "story = generator.get_story()\n",
        "script_text = render_story(story)\n",
        "print(script_text)\n",
        "\n",
        "# Render the prompts.\n",
        "prefix_text = render_prompts(generator.prompts)\n",
        "\n",
        "# Render the interventions.\n",
        "edits_text = ''\n",
        "for timestamp in sorted(generator.interventions):\n",
        "  edits_text += 'EDIT @ ' + str(timestamp) + '\\n'\n",
        "  edits_text += generator.interventions[timestamp] + '\\n\\n\\n'\n",
        "\n",
        "# Prepare the filenames for saving the story and prompts.\n",
        "timestamp_generation = datetime.datetime.now().strftime('%Y_%m_%d-%I_%M_%S_%p')\n",
        "title_ascii = re.sub('[^0-9a-zA-Z]+', '_',\n",
        "                     generator.title_str().strip()).lower()\n",
        "filename_script = f'{title_ascii}_{timestamp_generation}_script.txt'\n",
        "filename_prefix = f'{title_ascii}_{timestamp_generation}_prefix.txt'\n",
        "filename_edits = f'{title_ascii}_{timestamp_generation}_edits.txt'\n",
        "filename_config = f'{title_ascii}_{timestamp_generation}_config.json'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MPYk-Yp4EtW-"
      },
      "outputs": [],
      "source": [
        "#@title Save to Google Drive\n",
        "\n",
        "#@markdown Run this cell to save script, prefixes, edits and config as 4 text files on Google Drive under `dir_drive`.\n",
        "\n",
        "# Mount the Google Drive.\n",
        "import os\n",
        "import json\n",
        "from os import path\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive', force_remount=True)\n",
        "path_drive = '/content/gdrive/My Drive/'\n",
        "\n",
        "dir_drive = 'Dramatron_Outputs' #@param {type:\"string\"}\n",
        "config['file_dir'] = path_drive + dir_drive + \"/\"\n",
        "if not path.exists(config['file_dir']):\n",
        "  os.mkdir(config['file_dir'])\n",
        "  print(f\"Created directory: {config['file_dir']}\")\n",
        "\n",
        "print(f\"Outputs saved to: {config['file_dir']}\")\n",
        "\n",
        "# Save the script and prompt.\n",
        "path_script = config['file_dir'] + filename_script\n",
        "with open(path_script, 'wt') as f:\n",
        "  f.write(script_text)\n",
        "  print(f'Script saved to: {path_script}')\n",
        "\n",
        "path_prefix = config['file_dir'] + filename_prefix\n",
        "with open(path_prefix, 'wt') as f:\n",
        "  f.write(prefix_text)\n",
        "  print(f'Prefixes saved to: {path_prefix}')\n",
        "\n",
        "path_edits = config['file_dir'] + filename_edits\n",
        "with open(path_edits, 'wt') as f:\n",
        "  f.write(edits_text)\n",
        "  print(f'Edits saved to: {path_edits}')\n",
        "\n",
        "path_config = config['file_dir'] + filename_config\n",
        "with open(path_config, 'wt') as f:\n",
        "  config_clean = config.copy()\n",
        "  config_clean.pop('prefixes', None)\n",
        "  config_clean.pop('model_api_key', None)\n",
        "  json.dump(config_clean, f)\n",
        "  print(f'Config saved to: {path_config}')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "SiVzyJEZHyMs",
        "3OAW4PeE1z7N"
      ],
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
